2d数组的每列的第一个和最后一个数字1的索引

时间:2017-09-06 23:25:17

标签: python performance python-3.x numpy multidimensional-array

我有一个2d数组(Q),只包含0和1(二进制矩阵)。对于Q的每一列,我想找到值1出现的第一行和最后一行的索引。每列至少包含一个1

以下是一个例子:

[[1, 1, 1, 0, 0, 0, 0],
 [0, 1, 1, 1, 0, 0, 0],
 [1, 0, 0, 0, 1, 0, 1],
 [0, 0, 0, 1, 0, 1, 1],
 [1, 0, 1, 0, 0, 0, 0],
 [0, 0, 1, 0, 0, 0, 1],
 [0, 0, 0, 1, 0, 1, 0]]

boundsList = {0: (0, 4), 1: (0, 1), 2: (0, 5), 3: (1, 6), 4: (2, 2), 5: (3, 6), 6: (2, 5)}

我实现了一种算法,它可以工作,但对于大型数组来说效率不高:

boundsList = {}
for i in range (0, len(Q)):
    column = Q[:,i]
    indexesList = []
    for idx, pos in enumerate (column):
        if pos == 1:
            indexesList.append(idx)
    boundsList[i] = (indexesList[0], indexesList[-1])

任何人都可以建议另一个解决这个问题的简单方法吗?

2 个答案:

答案 0 :(得分:2)

让我们从您的阵列开始:

>>> Q
array([[1, 1, 1, 0, 0, 0, 0],
       [0, 1, 1, 1, 0, 0, 0],
       [1, 0, 0, 0, 1, 0, 1],
       [0, 0, 0, 1, 0, 1, 1],
       [1, 0, 1, 0, 0, 0, 0],
       [0, 0, 1, 0, 0, 0, 1],
       [0, 0, 0, 1, 0, 1, 0]])

获取包含1的第一行每列的索引:

>>> np.argmax(Q, axis=0) # Index of first appearance of 1
array([0, 0, 0, 1, 2, 3, 2])

获取包含1的最后一行的每列的索引:

>>> Q.shape[0] - np.argmax(Q[::-1, :], axis=0) - 1 # Index of last appearance of 1
array([4, 1, 5, 6, 2, 6, 5])

将它们组合成您喜欢的字典:

>>> dict(enumerate(zip( np.argmax(Q, axis=0), Q.shape[0] - np.argmax(Q[::-1, :], axis=0) - 1)))
{0: (0, 4), 1: (0, 1), 2: (0, 5), 3: (1, 6), 4: (2, 2), 5: (3, 6), 6: (2, 5)}

答案 1 :(得分:1)

可能最快的方法是使用argmax方法(因为它从两侧找到第一个最大值的位置),然后将其放入字典中。 argmax方法的开销(常数因子)远远小于使用np.argmax,因此,特别是对于小数组,方法会更快。

因为dictenumeratezip在列表上比数组快,我还将中间数组转换为列表(tolist方法是实现该目的的最快方法):

>>> dict(enumerate(zip(Q.argmax(axis=0).tolist(), 
...                    (Q.shape[0]-1-Q[::-1].argmax(axis=0)).tolist())))
{0: (0, 4), 1: (0, 1), 2: (0, 5), 3: (1, 6), 4: (2, 2), 5: (3, 6), 6: (2, 5)}

Q[::-1]是反向数组并且得到"不反转"索引我必须从Q.shape[0]-1中减去它们。