找到与模式匹配的numpy数组中所有条目的索引的最简洁方法是什么?

时间:2014-02-24 14:25:27

标签: python numpy boolean

有没有办法整齐地提取匹配模式的数组中所有值的索引?

以下示例生成一个元组列表,其索引与数组A的元素匹配为True

import numpy as np

A = np.array([[False, False, False, True],
              [False, True, False, False],
              [False, False, False, False]])
indices = list()
for i, K in enumerate(A):
    for j, L in enumerate(K):
        if L:
            indices.append((i, j))

这很有效,但看起来很丑陋而且“unpythonic”。这个解决方案的一个具体缺陷是它依赖于A的维数。有更好的方法吗?如果将索引生成为元组列表或其他易于访问的内容,则无关紧要。

1 个答案:

答案 0 :(得分:2)

您可以使用np.argwhere

In [1]: A = np.array([[False, False, False, True],
   ...:               [False, True, False, False],
   ...:               [False, False, False, False]])

In [2]: np.argwhere(A == True)
Out[2]: 
array([[0, 3],
       [1, 1]])