python - 如何返回 numpy 数组的列主迭代器？

Question

numpy 中的

ndarray 对象有一个 flat 属性(例如 array.flat)，允许迭代通过它的元素。例如:

>>> x = np.arange(1, 7).reshape(2, 3)
>>> x
array([[1, 2, 3],
       [4, 5, 6]])
>>> x.flat[3]
4

但是我怎样才能返回列优先的一维迭代器，以便上面的示例返回 5 而不是 4？

Answer 1

方法#1

您可以使用 .ravel('F') 对列进行主要排序，然后进行索引 -

x.ravel('F')[3]

sample 运行-

In [100]: x
Out[100]: 
array([[1, 2, 3],
       [4, 5, 6]])

In [101]: x.ravel('F')[3]
Out[101]: 5

这将在选择元素之前创建整个数组的副本 -

In [161]: np.may_share_memory(x, x.ravel())
Out[161]: True

In [162]: np.may_share_memory(x, x.ravel('F'))
Out[162]: False

因此这可能不是内存效率最高的一个。为了更好，让我们转向另一种方法。

---

方法 #2

我们可以从列优先的有序索引中获取行索引和列索引，然后用它简单地索引到数组中 -

x[np.unravel_index(3, np.array(x.shape)[::-1])]

sample 运行-

In [147]: x
Out[147]: 
array([[1, 2, 3],
       [4, 5, 6]])

In [148]: idx = np.unravel_index(3, np.array(x.shape)[::-1])

In [149]: idx
Out[149]: (1, 1) # row, col indices obtained in C order

In [150]: x[idx]
Out[150]: 5

这里没有任何复制、展平或拼接，仅使用索引，因此在内存和性能方面都应该是高效的。