python - 如何从列表/ndarray 中获取索引？

Question

我有一个列表，如下所示:

[[0,1,2], [1,2,3], [2,3,4], [3,4,5]]

我可以把它变成这样的数组:

array([[0,1,2],
       [1,2,3],
       [2,3,4],
       [3,4,5]])

所以我总共有 4 行，每行有 3 列。现在我想找到所有大于 2 的元素的索引，所以对于整个矩阵，索引应该是:

((1,2),(2,1),(2,2),(3,1),(3,2),(3,3))

然后对于每一行，我将随机选择一个指示大于 2 的值的 col 索引。现在我的代码是这样的:

a = np.array([[0,1,2],[1,2,3],[2,3,4],[3,4,5]]
out = np.ones(4)*-1
cur_row = 0
col_list = []
for r,c in np.nonzero(a>2):
    if r == cur_row:
        col_list.append(c)
    else:
        cur_row = r
        shuffled_list = shuffle(col_list)
        out[r-1] = shuffled_list[0]
        col_list = []
        col_list.append(c) 

我希望得到一个看起来像这样的输出:

array([-1, 2, 1, 2])

但是，现在当我运行我的代码时，它显示

ValueError: too many values to unpack

有人知道我如何解决这个问题吗？或者我应该怎么做才能达到我的目标？我只是想尽可能快地运行代码，所以任何其他好的想法也非常受欢迎。

Answer 1

试试这个。

import numpy as np

arr = np.array([[0,1,2],
               [1,2,3],
               [2,3,4],
               [3,4,5]])
indices = np.where(arr>2)

for r, c in zip(*indices):
    print(r, c)

打印

1 2
2 1
2 2
3 0
3 1
3 2

所以，它应该可以工作。您也可以使用 itertools.izip，在这种情况下它甚至是更好的选择。

纯 numpy 解决方案(感谢@AshwiniChaudhary 的提议):

for r, c in np.vstack(np.where(arr>2)).T:
    ...

虽然我不确定这会比使用 izip 或 zip 更快。