python - 来自形状的 Numpy 广播索引

Question

我有 2 个数组形状，它们可以相互广播。

例如(2, 2, 1) 和 (2, 3)

我想要一个采用这些形状并给我一个迭代器的函数，该迭代器返回这些数组中的索引以及将一起广播的这些形状，以及生成的输出数组中的索引。

iter, output_shape = broadcast_indeces_iterator((2, 2, 1), (2, 3))
assert output_shape == (2, 2, 3)
for in1_ix, in_2_ix, out_ix in iter:
    print (in1_ix, in_2_ix, out_ix)    

输出结果:

(0, 0, 0), (0, 0), (0, 0, 0)
(0, 0, 0), (0, 1), (0, 0, 1)
(0, 0, 0), (0, 2), (0, 0, 2)
(0, 1, 0), (1, 0), (0, 1, 0)
(0, 1, 0), (1, 1), (0, 1, 1)
(0, 1, 0), (1, 2), (0, 1, 2)
(1, 0, 0), (0, 0), (1, 0, 0)
(1, 0, 0), (0, 1), (1, 0, 1)
(1, 0, 0), (0, 2), (1, 0, 2)
(1, 1, 0), (1, 0), (1, 1, 0)
(1, 1, 0), (1, 1), (1, 1, 1)
(1, 1, 0), (1, 2), (1, 1, 2)

np.broadcast 做了一些接近但需要实际创建的数组。

• NumPy 用户请注意:如果 np.broadcast 有一个额外的参数允许您不迭代，例如，在最后 2 个维度上，那就太好了。这也可以解决我的问题。

Answer 1

import numpy as np
x = 10*np.arange(4).reshape((2, 2, 1))
y = 100*np.arange(6).reshape((2, 3))

z = np.nditer([x, y], flags=['multi_index', 'c_index'], order='C')
for a,b in z:
    print(np.unravel_index(z.index % x.size, x.shape)
          , np.unravel_index(z.index % y.size, y.shape)
          , z.multi_index)

产量

((0, 0, 0), (0, 0), (0, 0, 0))
((0, 1, 0), (0, 1), (0, 0, 1))
((1, 0, 0), (0, 2), (0, 0, 2))
((1, 1, 0), (1, 0), (0, 1, 0))
((0, 0, 0), (1, 1), (0, 1, 1))
((0, 1, 0), (1, 2), (0, 1, 2))
((1, 0, 0), (0, 0), (1, 0, 0))
((1, 1, 0), (0, 1), (1, 0, 1))
((0, 0, 0), (0, 2), (1, 0, 2))
((0, 1, 0), (1, 0), (1, 1, 0))
((1, 0, 0), (1, 1), (1, 1, 1))
((1, 1, 0), (1, 2), (1, 1, 2))