对于一维 numpy 数组中的每个元素，从第二个数组中找到最低索引元素，使得值和索引都大于第一个

Question

我有两个长度约为 100 万的一维 numpy 数组。对于第一个数组中的每个元素 x，我想要第二个数组中的最低索引元素 y，使得 y > x 且 arg(y) > arg(x)，即 y 的索引 > x 的索引。

假设，

arr1 = [3,2,1,0,3,2,3] arr2 = [0,2,1,2,3,4,6,5]

输出是-

值 = [4,3,2,3,4,6,5] 索引 = [5,4,3,4,5,6,7]

完全迷失了。无法使用广播，因为数组的长度为 100 万。我尝试过广播来获取 y > x 的元素，但我无法弄清楚如何应用第二个条件。

Answer 1

arr2

可以广播到

arr1

以通过值（

>

）和索引（arr2 > arr1[:, None]）执行

np.indices(arr2.shape)[0] > np.indices(arr1.shape)[0][:, None]

（

大于

）比较。
这涉及在

arr1

上放置一个新轴，使其成为列向量，并最终创建 2 个布尔掩码。
这些掩码的交集允许找到与

np.argmax

的第一个有效索引，然后将其用于用

np.take

取值。
整个解决方案如下：

indices = (np.argmax((arr2 > arr1[:, None]) 
           & (np.indices(arr2.shape)[0] > np.indices(arr1.shape)[0][:, None]), axis=1))
values = np.take(arr2, indices)

---

indices
array([5, 4, 3, 4, 5, 6, 7])

values
array([4, 3, 2, 3, 4, 6, 5])