python - 使用索引列表有效地从数组中提取值

Question

给定一个二维 NumPy 数组 a 和存储在 index 中的索引列表，必须有一种非常有效地提取列表值的方法。如下使用 for 循环大约需要 5 毫秒，这对于提取 2000 个元素来说似乎非常慢:

import numpy as np
import time

# generate dummy array 
a = np.arange(4000).reshape(1000, 4) 
# generate dummy list of indices
r1 = np.random.randint(1000, size=2000)
r2 = np.random.randint(3, size=2000)
index = np.concatenate([[r1], [r2]]).T

start = time.time()
result = [a[i, j] for [i, j] in index]
print time.time() - start

如何提高提取速度？ np.take 在这里似乎不合适，因为它会返回一个二维数组而不是一维数组。

Answer 1

您可以使用 advanced indexing这基本上意味着从 index 数组中提取行和列索引，然后使用它从 a 中提取值，即 a[index[:,0]，索引[:,1]] -

%timeit a[index[:,0], index[:,1]]
# 12.1 µs ± 368 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit [a[i, j] for [i, j] in index]
# 2.22 ms ± 105 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)