使用 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 毫秒