通过其他数组作为索引来索引 numpy 数组

Question

我正在尝试获取数组

a = [1,5,4,5,7,8,9,8,4,13,43,42]

和数组

b = [3,5,6,2,7]

我希望 b 成为 a 中的索引，例如一个新数组是

[a[b[0]], a[b[1]], a[b[2]], a[b[3]] ...]

因此 b 中的值是 a 的索引。 a 中有 500k 个条目，b 中有 500k 个条目（大约）。 有没有一种快速的方法可以在 numpy 中启动所有核心来做到这一点？ 我已经在 for 循环中做得很好了，它是 sloooooooowwwwww。

编辑以澄清。 该解决方案必须适用于 2D 和 3D 阵列。 所以也许

b = [(2,3), (5,4), (1,2), (1,0)]

我们想要

c = [a[b[0], a[b[1], ...]

Answer 1

并不是说它很快，但 numpy 的方法就是：

a[b]

输出：

数组([5,8,9,4,8])

Answer 2

这可以在 NumPy 中使用高级索引来完成。正如 Christian 的回答 所指出的，在一维情况下，你只需写：

a[b]

这相当于：

[a[b[x]] for x in range(b.shape[0])]

但是，在高维情况下，您需要为索引的每个维度都有单独的列表。这意味着，你不能这样做：

a = np.random.randn(7, 8, 9)  # 3D array
b = [(2, 3, 0), (5, 4, 1), (1, 2, 2), (1, 0, 3)]
print(a[b])  # this is incorrect

但你可以这样做：

b0, b1, b2 = zip(*b)
print(a[b0, b1, b2])

您还可以使用

np.take

:

print(np.take(a, b))

Answer 3

我通过编写 numpy 的 C 扩展（称为张量加权插值传输）来解决这个问题，以获得速度和多线程。在纯 Python 中，每 200x100x3 图像缩放和淡入淡出需要 3 秒，而在具有 8 核的多线程 C 中，相同操作需要 0.5 毫秒。

核心 C 代码最终就像这样

t2[dstidxs2[i2] + doff1] += t1[srcidxs2[i2] + soff1] * w1 * ws2[i2];

doff1 是目标数组中的偏移量等。w1 和 ws2 是插值权重。 所有代码都在 C 语言中进行了超优化，以提高速度。 （不是代码大小或可维护性）

所有代码均可在 https://github.com/RMKeene/twit 和 PyPI 上获取。

我预计未来会进一步优化，例如所有权重均为 1.0 的特殊情况。

--- 2024 年 12 月注：我放弃了 Twit，因为它不再是我的人工智能研究所需要的。