通过比较数组之间的元素位置来获取numpy数组元素的索引

Question

上下文

我在numpy中有以下示例阵列：

import numpy as np
# All arrays in this example have the shape (15,)
# Note: All values > 0 are unqiue!
a         = np.array([8,5,4,-1,-1, 7,-1,-1,12,11,-1,-1,14,-1,-1])
reference = np.array([0,1,2, 3, 4, 5, 6, 7, 8, 9,10,11,12,13,14])
lookup    = np.array([3,6,0,-2,-2,24,-2,-2,24,48,-2,-2,84,-2,-2])

我的目标是在a中找到reference内部的元素，然后在a中获得索引，并使用它来提取lookup中的相应元素。

找出匹配元素及其索引与np.flatnonzero( np.isin() )一起工作。我还可以查找相应的值：

# Example how to find the index
np.flatnonzero( np.isin( reference, a) )
# -> array([ 4,  5,  7,  8, 11, 12, 14])
# Example how to find corresponding values:
lookup[ np.flatnonzero( np.isin( a, reference) ) ]
# -> array([ 3,  6,  0, 24, 24, 48, 84], dtype=int64)

问题

我想用我在引用之后查找的值填充数组z。这意味着，例如z的第8个元素对应于reference中的第8元素的查找值中的第八个元素(＝8)。该值将为3(reference[8] -> a[0] because a==8 here -> lookup[0] -> 3)。

z = np.zeros(reference.size)
z[np.flatnonzero(np.isin(reference, a))] = ? -> numpy-array of correctly ordered lookup_values

z的预期结果为：

z = [ 0  0  0  0  0  6  0  24  3  0  0 48 24 0  84]

我无法理解这一点；由于性能原因，我不得不避免使用for-loops，并且需要一个纯numpy解决方案(最好没有udfs)。

如何根据正确位置的查找值填充z

注意：如上面代码中所述，所有值a > 0都是唯一的。因此，不需要考虑a < 0的重复值

Answer 1

您说"由于性能原因，必须避免循环"，所以我假设您的真实世界数据结构a将很大(数千或数百万个元素？)。由于np.isin(reference, a)在a中对reference的每个元素执行线性搜索，因此您的运行时将为O(len(reference) * len(a))。

我强烈建议对a使用dict，允许在O(1)中查找reference的每个元素，并使用for在python中循环。对于足够大的a，这将优于np.isin执行的"快速"线性搜索。

Answer 2

我能想到的最自然的方法是将a和lookup视为字典：

In [82]: d = dict(zip(a, lookup))
In [83]: np.array([d.get(i, 0) for i in reference])
Out[83]: array([ 0,  0,  0,  0,  0,  6,  0, 24,  3,  0,  0, 48, 24,  0, 84])

这确实有一点内存开销，但如果reference不是太大，那就没什么疯狂的了。

Answer 3

我其实有一个启示。

# Initialize the result
# All non-indexed entries shall be 0
z = np.zeros(reference.size, dtype=np.int64)

现在评估a中哪些元素是相关的：

mask = np.flatnonzero(np.isin(a, reference))
# Short note: If we know that any positive element of a is a number
# Which has to be in the reference, we can also shorten this to
# a simple boolean mask. This will be significantly faster to process.
mask = (a > 0)

现在使用以下技巧：所有值a > 0都是唯一的。此外，它们的值与reference中的位置相对应(例如，a中的8应与reference中的第8个位置相对应。因此，我们可以将这些值本身用作索引：

z[ a[mask] ] = lookup[mask]

这导致了所需的结果：

z = [ 0  0  0  0  0  6  0  24  3  0  0 48 24 0  84]