numpy数组中多个元素的索引

Posted 2023-04-18

【中文标题】numpy数组中多个元素的索引

【英文标题】：Indices of multiple elements in a numpy array

【发布时间】：2018-07-22 14:54:43

【问题描述】：

我有一个numpy数组和一个列表如下

y=np.array([[1],[2],[1],[3],[1],[3],[2],[2]])
x=[1,2,3]

我想返回一个数组元组，每个数组都包含 y 中 x 的每个元素的索引。 即

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

这是否可以以矢量化方式完成（没有任何循环）？

【问题讨论】：

来自x的所有元素都保证在y中吗？

是的，x 中的所有元素都保证在y 中

只是看看所涉及的参数 - 在您的实际使用案例中，y 和 x 的典型长度是多少？

由于结果列表的大小可能不同，因此不太可能存在“矢量化”解决方案。

【参考方案1】：

一个解决方案是map

y = y.reshape(1,len(y))
map(lambda k: np.where(y==k)[-1], x)

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

---

合理的表现。对于 100000 行，

%timeit list(map(lambda k: np.where(y==k), x))
3.1 ms ± 113 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

【参考方案2】：

尝试以下方法：

y = y.flatten()
[np.where(y == searchval)[0] for searchval in x]

【讨论】：

无法避免for循环，可以吗？

不，这真的不可能。不过这应该很快。

map 无论如何都使用内部 for 循环，所以我怀疑这更快

map 和列表理解之间没有显着的速度差异。两者都不会将对 where 的调用转移到 C 代码中。

【参考方案3】：

您可以使用collections.defaultdict 后跟理解：

y = np.array([[1],[2],[1],[3],[1],[3],[2],[2]])
x = [1,2,3]

from collections import defaultdict

d = defaultdict(list)
for idx, item in enumerate(y.flat):
    d[item].append(idx)

res = tuple(np.array(d[k]) for k in x)

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

【讨论】：

我第一次使用 defaultdict，但意识到使用已知的 x 列表，普通字典也能正常工作，而且可能更快。

@hpaulj，我相信你是对的。我对您的解决方案的唯一问题是 dd.values() 仅在 Python 3.6+ 中订购。

【参考方案4】：

对于这个小例子，字典方法实际上更快（然后是`wheres）：

dd = i:[] for i in [1,2,3]
for i,v in enumerate(y):
   v=v[0]
   if v in dd:
       dd[v].append(i)
list(dd.values())

这个问题已经出现在其他 SO 问题中。已经提出了使用 unique 和 sort 的替代方案，但它们更复杂，更难重新创建 - 而且不一定更快。

对于numpy，这不是一个理想的问题。结果是一个数组列表或不同大小的列表，这是一个很好的线索，表明简单的“矢量化”全数组解决方案是不可能的。如果速度是一个足够重要的问题，您可能需要查看 numba 或 cython 实现。

不同的方法可能具有不同的相对时间，具体取决于值的组合。很少有唯一值，但较长的子列表可能偏爱使用重复 where 的方法。许多具有短子列表的唯一值可能会支持在 y 上进行迭代的方法。