在数组中搜索所有匹配项并返回匹配项的索引

Question

我试图创建一个数组，其中包含一个（非常大的数组）与一组唯一值匹配的所有行。问题是，大数组将有多个匹配的行，我需要将它们全部存储在新数组的同一行中

使用for循环遍历每个唯一的值是可行的，但是速度太慢而不可用。我一直在寻找一个矢量化的解决方案，但没有成功。任何帮助都将不胜感激

    arrStart = []
    startRavel = startInforce['pol_id'].ravel()
    for policy in unique_policies:
        arrStart.append(np.argwhere(startRavel == policy))

新数组将具有与unique values数组相同的长度，但每个元素将是一个列表，其中包含与大数组中的唯一值匹配的所有行

示例输入如下： 星际旅行=[1,2,2,2,3,3] 唯一的\u策略=[1,2,3]

输出： arrStart=[[0]，[1,2,3]，[4,5]]

Answer 1

NumPy的一个可能选项，类似于您的列表，但在列表理解中变平：

startRavel = np.array([1,2,2,2,3,3])
unique_policies = np.array([1,2,3])

[np.argwhere(startRavel == policy).flatten() for policy in unique_policies]
#=> [array([0]), array([1, 2, 3]), array([4, 5])]

---

或者，使用^{}：

[np.flatnonzero(startRavel == policy) for policy in unique_policies]

---

发电机版本：

def matches_indexes(startRavel, unique_policies):
  for policy in unique_policies:
    yield np.flatnonzero(startRavel == policy)