迭代 NumPy 数组的列和另一个数组的元素？答案

【问题标题】：Iterate over columns of a NumPy array and elements of another one?迭代 NumPy 数组的列和另一个数组的元素？
【发布时间】：2014-09-21 16:51:18
【问题描述】：

我正在尝试复制zip(a, b) 的行为，以便能够同时沿两个NumPy 数组循环。特别是，我有两个数组a 和b：

a.shape=(n,m)
b.shape=(m,)

我想为每个循环获取一个 a 列和一个 b 元素。

到目前为止，我已经尝试了以下方法：

for a_column, b_element in np.nditer([a, b]):
    print(a_column)

但是，我打印的是元素a[0,0]，而不是我想要的列a[0,:]。

我该如何解决这个问题？

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

您仍然可以在 numpy 数组上使用 zip，因为它们是可迭代的。

在您的情况下，您需要先转置a，使其成为形状为(m,n) 的数组，即长度为m 的可迭代：

for a_column, b_element in zip(a.T, b):
    ...

【讨论】：

这是最pythonic的答案。 np.nditer 可能更强大，但更丑陋。例如，您必须回忆一下 C 排序和 Fortran 排序之间的区别。

【解决方案2】：

在shallow iteration with nditer 中调整我的答案， nditer 和 ndindex 可用于通过生成索引来迭代行或列。

In [19]: n,m=3,4
In [20]: a=np.arange(n*m).reshape(n,m)
In [21]: b=np.arange(m)

In [22]: it=np.nditer(b)
In [23]: for i in it: print a[:,i],b[i]
[0 4 8] 0
[1 5 9] 1
[ 2  6 10] 2
[ 3  7 11] 3

In [24]: for i in np.ndindex(m):print a[:,i],b[i]
[[0]
 [4]
 [8]] 0
[[1]
 [5]
 [9]] 1
[[ 2]
 [ 6]
 [10]] 2
[[ 3]
 [ 7]
 [11]] 3
In [25]:

ndindex 使用如下迭代器：it = np.nditer(b, flags=['multi_index']。

对于像这样的单一维度上的迭代，for i in range(m): 也同样有效。

同样来自另一个线程，这里有一个使用order 在没有索引的情况下进行迭代的技巧：

In [28]: for i,j in np.nditer([a,b],order='F',flags=['external_loop']):
    print i,j
[0 4 8] [0 0 0]
[1 5 9] [1 1 1]
[ 2  6 10] [2 2 2]
[ 3  7 11] [3 3 3]

【讨论】：

【解决方案3】：

通常，由于 NumPy 具有broadcast arrays 的能力，没有必要逐个遍历数组的列。例如，如果a 的形状为(n,m)，b 的形状为(m,)，那么您可以添加a+b，b 将自动广播自己以形成(n, m)。

此外，如果可以通过对整个数组 a 的操作而不是通过使用 Python for-loop 对 a 的片段（例如列）进行操作来表达，您的计算将完成得更快。

话虽如此，遍历a 的列的最简单方法是遍历索引：

for i in np.arange(b.shape[0]):
    a_column, b_element = a[:, i], b[i]
    print(a_column)

【讨论】：