Numpy：对于一个数组中的每个元素，在另一个数组中找到其索引位置。

Question

我有两个一维数组x和y，其中一个比另一个小。我想找到y中每个元素在x中的索引。

我发现了两种朴素的方法来实现这个目标，第一种方式速度较慢，第二种方式占用内存较高。

较慢的方式

indices= []
for iy in y:
    indices += np.where(x==iy)[0][0]

内存占用过多

xe = np.outer([1,]*len(x), y)
ye = np.outer(x, [1,]*len(y))
junk, indices = np.where(np.equal(xe, ye))

有没有更快或者内存占用更少的方法？理想情况下，搜索应该利用这样一个事实：我们在一个列表中搜索的不是一个，而是很多东西，因此稍微更容易并行化。如果你不假设 y 的每个元素都在 x 中，那就更好了。

- Chris

Answer 1

我想提出一种简洁的解决方案：

indices = np.where(np.in1d(x, y))[0]

结果是一个数组，其中包含x数组的索引，这些索引对应于在x中找到的y元素。

如果需要，可以在没有numpy.where的情况下使用它。

- RomanS

---

这应该是最佳答案。即使x的值重复或不存在，它也能正常工作。涉及searchsorted的答案很复杂、奇怪、不自然。 - Wilmer E. Henao

虽然这个方法可以返回y中存在于x中的元素的索引，但返回的索引顺序与x中的值的顺序不匹配。例如：x=np.array([1,2,3,4,5]); y=np.array([5,4,3,2,1])。上述方法返回array([0,1,2,3,4])，因此x[0]=1与y[0]=5相匹配，而这并不是想要的结果... - ccbunney

"in1d() solutions just do not work. Take y = np.array([10, 5, 5, 1, 'auto', 6, 'auto', 1, 5, 10, 10, 'auto']) and x = np.array(['auto', 5, 6, 10, 1]). You would expect [3, 1, 1, 4, 0, 2, 0, 4, 3, 3, 0]. np.where(np.in1d(x, y))[0] doesn't yield that."使用in1d()函数解决问题无效。以y = np.array([10, 5, 5, 1, 'auto', 6, 'auto', 1, 5, 10, 10, 'auto']) 和x = np.array(['auto', 5, 6, 10, 1])为例，你期望得到的结果是[3, 1, 1, 4, 0, 2, 0, 4, 3, 3, 0]。但是np.where(np.in1d(x, y))[0]并不能输出这个结果。 - hermidalc

这个代码片段的作用是判断 x 中的元素是否存在于 y 中，并返回它们在 x 中对应的索引值。但它并不会给出每个在 x 中的元素相对应在 y 中的索引值。 - Brian Pollack

Answer 2

正如Joe Kington所说，searchsorted()可以非常快速地搜索元素。要处理不在x中的元素，您可以使用原始y检查搜索结果，并创建掩码数组：

import numpy as np
x = np.array([3,5,7,1,9,8,6,6])
y = np.array([2,1,5,10,100,6])

index = np.argsort(x)
sorted_x = x[index]
sorted_index = np.searchsorted(sorted_x, y)

yindex = np.take(index, sorted_index, mode="clip")
mask = x[yindex] != y

result = np.ma.array(yindex, mask=mask)
print result

结果是：

[-- 3 1 -- -- 6]

- HYRY

Answer 3

这个怎么样？

它假设y的每个元素都在x中（即使对于不在x中的元素也会返回结果！），但它要快得多。

import numpy as np

# Generate some example data...
x = np.arange(1000)
np.random.shuffle(x)
y = np.arange(100)

# Actually preform the operation...
xsorted = np.argsort(x)
ypos = np.searchsorted(x[xsorted], y)
indices = xsorted[ypos]

- Joe Kington

---

太棒了。确实快多了。我正在添加 assert na.all(na.intersect1d(x,y) == na.sort(y)) 来限制输入，以便 y 是 x 的子集。谢谢！ - Chris

如果y是x的子集，那么这个方法就会起作用。否则会引发IndexError。 - undefined

Answer 4

我认为下面这个版本更加清晰:

np.where(y.reshape(y.size, 1) == x)[1]

比起 indices = np.where(y[:, None] == x[None, :])[1] ，你不需要将x广播成2D数组。我发现这种解决方案是最好的，因为它不像搜索排序(searchsorted())或基于in1d()的解决方案那样，可以处理重复项并且不在乎任何东西是否已排序。这对我很重要，因为我想要x按照特定的自定义顺序。

- hermidalc

---

更清晰并不意味着更低效。 - Mad Physicist

我猜你可以进一步简化为 y.reshape(-1, 1)。 - Roman Zh.

实际上，np.where(y[:, None] == x)[1] 就足够了。 - Dmitriy Work

Answer 5

我会这样做：

indices = np.where(y[:, None] == x[None, :])[1]

与您的内存消耗方式不同，此方法利用广播直接生成二维布尔数组，而不需要为x和y分别创建二维数组。

- Jun Saito

---

记录一下，这也占用了内存。 - romeric

是的，我的意思是它占用的内存更少。我认为我的版本在保持代码清晰的同时占用更少的内存方面达成了很好的折衷。 - Jun Saito

这种方法的速度比被接受的答案慢了1000倍。 - Alex Kaszynski

Answer 6

numpy_indexed包含一个函数可以完成这个功能（免责声明：我是它的作者）：

import numpy_indexed as npi
indices = npi.indices(x, y, missing='mask')

目前，如果y中的所有元素在x中都不存在，这段代码将会抛出KeyError错误。但是我应该添加kwarg参数，使用户可以选择用-1或其他标记标记这些项。

这段代码的效率应该与当前接受的答案相同，因为实现方式类似。不过，numpy_indexed更加灵活，还允许搜索多维数组的行索引。

编辑：我已经改变了缺失值的处理方式；'missing' kwarg现在可以设置为'raise'、'ignore'或'mask'。在后一种情况下，您将得到一个与y长度相同的掩码数组，您可以调用.compressed()来获取有效的索引。请注意，如果这是您想知道的所有内容，则还有npi.contains(x, y)。

- Eelco Hoogendoorn

Answer 7

另一个解决方案是：

a = np.array(['Bob', 'Alice', 'John', 'Jack', 'Brian', 'Dylan',])
z = ['Bob', 'Brian', 'John']
for i in z:
    print(np.argwhere(i==a))

- NSVR

Answer 8

使用以下代码行：-

indices = np.where(y[:, None] == x[None, :])[1]

- Kaushal Gupta

Answer 9

我的解决方案可以处理多维的x。默认情况下，它将返回一个与x形状相对应的标准numpy数组中相应y索引的数组。

如果不能假设y是x的子集，则设置masked=True以返回一个掩码数组（这会有性能损失）。否则，您仍将获得未包含在y中的元素的索引，但它们可能对您没有用处。

HYRY和Joe Kington的答案对此很有帮助。

# For each element of ndarray x, return index of corresponding element in 1d array y
# If y contains duplicates, the index of the last duplicate is returned
# Optionally, mask indices where the x element does not exist in y

def matched_indices(x, y, masked=False):
    # Flattened x
    x_flat = x.ravel()

    # Indices to sort y
    y_argsort = y.argsort()

    # Indices in sorted y of corresponding x elements, flat
    x_in_y_sort_flat = y.searchsorted(x_flat, sorter=y_argsort)

    # Indices in y of corresponding x elements, flat
    x_in_y_flat = y_argsort[x_in_y_sort_flat]

    if not masked:
        # Reshape to shape of x
        return x_in_y_flat.reshape(x.shape)

    else:
        # Check for inequality at each y index to mask invalid indices
        mask = x_flat != y[x_in_y_flat]
        # Reshape to shape of x
        return np.ma.array(x_in_y_flat.reshape(x.shape), mask=mask.reshape(x.shape))

- Stefan

Answer 10

更紧凑的解决方案：

indices, = np.in1d(a, b).nonzero()

- MathPass