从np.tril_indices“反转”索引的有效方法

时间:2018-06-17 14:52:20

标签: python arrays numpy indexing

np.tril_indices(3)返回大小为3的方阵的下三角索引:

(array([0, 1, 1, 2, 2, 2], dtype=int64),
array([0, 0, 1, 0, 1, 2], dtype=int64))

获得最有效的方法是什么:

(array([0, 1, 1, 2, 2, 2], dtype=int64),
array([0, 1, 0, 2, 1, 0], dtype=int64))

感谢。

2 个答案:

答案 0 :(得分:4)

一种方法是

>>> np.subtract.accumulate(np.tril_indices(3), 0)
array([[0, 1, 1, 2, 2, 2],
       [0, 1, 0, 2, 1, 0]])

答案 1 :(得分:1)

此处有cumsummaximum.accumulate -

def reverse_tril_indices(n):
    r = np.arange(n)
    rr = np.arange(n,0,-1)
    N = n*(n+1)//2

    shifts1 = rr[:-1].cumsum()
    shifts2 = r.cumsum()

    id_arr = np.ones(N,dtype=int)    
    id_arr[shifts1] = -rr[1:]
    id_arr[0] = 0
    id1 = id_arr.cumsum()[::-1]

    id2_arr = np.zeros(N,dtype=int) # use dtype=np.uint16 for further boost
    id2_arr[shifts2] = r
    id2 = np.maximum.accumulate(id2_arr)
    return id2, id1

示例运行 -

In [75]: reverse_tril_indices(3)
Out[75]: (array([0, 1, 1, 2, 2, 2]), array([0, 1, 0, 2, 1, 0]))

In [77]: reverse_tril_indices(5)
Out[77]: 
(array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 4]),
 array([0, 1, 0, 2, 1, 0, 3, 2, 1, 0, 4, 3, 2, 1, 0]))

n=5k上的时间 -

# @Paul Panzer's soln
In [83]: %timeit np.subtract.accumulate(np.tril_indices(5000), 0)
1 loop, best of 3: 278 ms per loop

In [84]: %timeit reverse_tril_indices(5000)
10 loops, best of 3: 83.4 ms per loop