使用 arrays 的列表索引二维数组

Question

我想根据 arrays 的列表索引一个二维数组。

a = np.array([
    [1,2,3],
    [4,5,6]])

idx = [np.array([0,1], np.array([0,2])]

然后我想要的是 idx 中的第一个元素应该给出a[0,1]和第二个a[0,2]等等这样：

a[fixed_idx] = array([2,3])

Answer 1

IIUC，你可以这样做：

a[tuple(zip(*idx))]

output: array([2, 3])

Answer 2

假设您有更多索引，例如：

dummy_idx = [n for n in np.random.randint(100, size=(1000000, 2))]

然后你需要获得高级索引x和y使得a[x, y]给出你所期望的。

有两种简单的方法可以做到这一点：

1. x, y = zip(*dummy_idx)
2. x, y = np.transpose(dummy_idx)

第一种方法相当慢，因为numpy arrays 不是为快速迭代而设计的，因此与numpy量化操作相比，访问它们的项目需要相当长的时间。 另一方面， np.transpose将多个 arrays 收集到一个新数组中，这更糟糕，因为每个步骤都需要将它们保存在这个新数组的某个地方，这更昂贵。

这是一个危险信号，您正在尝试使用数据结构 numpy 不是为之设计的。 实际上，如果您使用大量小的 arrays，它会很慢。

但是，有两种方法np.ndarray.tolist和np.ndarray.tobytes针对重复使用进行了优化。 所以你可以利用这个优势并尝试以 30% 的速度模仿np.transpose(dummy_idx)的行为：

ls = []
for n in dummy_idx: 
    ls.extend(n.tolist())
x, y = np.fromiter(ls, dtype=dummy_idx[0].dtype).reshape(-1, 2).T

和

b = bytearray()
for n in dummy_idx: 
    b.extend(n.tobytes())
x, y = np.frombuffer(b, dtype=dummy_idx[0].dtype).reshape(-1, 2).T

结果

• zip - 161 毫秒
• np.transpose - 205 毫秒
• np.fromiter - 117 毫秒
• np.frombuffer - 117 毫秒
• 单循环dummy_idx （相比之下）- 16 毫秒