如何从 list/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))

然后对于每一行，我会随机挑选一个col索引，它表示一个值大于2。现在我的代码是这样的：

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)

版画

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

一个纯粹的 numpy 解决方案（感谢@AshwiniChaudhary 的提议）：

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

尽管我不确定这会比使用 izip 或 zip 更快。

Answer 2

您可以将数组与您的值进行比较并使用 where。

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

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

获取元组

list(zip(*np.where(a>2)))

[(1, 2), (2, 1), (2, 2), (3, 0), (3, 1), (3, 2)]

Answer 3

我已经弄出来了，代码应该是：

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 zip(*(np.nonzero(a>2))):
    if  r == cur_row:
        col_list.append(c)
    else:
        cur_row = r
        shuffle(col_list)
        if len(col_list) == 0:
            out[r-1] = -1
        else:
            out[r-1] = col_list[0]
        col_list = []
        col_list.append(c)

shuffle(col_list)
if len(col_list) == 0:
    out[len(out)-1] = -1
else:
    out[len(out)-1] = col_list[0]

最后但在 forloop 之外的部分是确保处理最后一行。

它适用于我的情况。

如何从 list/ndarray 获取索引？

How to get indices from a list/ndarray?

python

numpy

matrix

indices