python - Iterating through multiple numpy arrays simultaneously

Question

I have a numpy array with shape (n,), for example:

[
    [0, 1],         # A
    [0, 1, 2, 3],   # B
    [0, 1, 2]       # C
]

How can I iterate through all 'tuples' formed by this array, in this case I would expect the return to be (0,0,0),(1,0,0),(0,1,0),(0,2,0),(0,3,0),(1,1,0)... and so on

标签： pythonnumpy

Answer 1

这是适合您的 numpy 解决方案。

inp = [
    [0, 1,],        
    [0, 1, 2, 3],  
    [0, 1, 2]       
]

现在您可以简单地

import numpy as np

shp = []
for sub_list in inp:
    shp.append(len(sub_list))

arr = np.ones(shp)

result = np.where(arr)

tuples = [t for t in zip(*result)]

出去：

[(0, 0, 0),
 (0, 0, 1),
 (0, 0, 2),
 (0, 1, 0),
 (0, 1, 1),
 (0, 1, 2),
 (0, 2, 0),
 (0, 2, 1),
 (0, 2, 2),
 (0, 3, 0),
 (0, 3, 1),
 (0, 3, 2),
 (1, 0, 0),
 (1, 0, 1),
 (1, 0, 2),
 (1, 1, 0),
 (1, 1, 1),
 (1, 1, 2),
 (1, 2, 0),
 (1, 2, 1),
 (1, 2, 2),
 (1, 3, 0),
 (1, 3, 1),
 (1, 3, 2)]

它的工作方式是构建一个数组，其维度是列表的长度。
然后你得到这个数组的多维索引，这恰好是你想要的。

如果您还想访问相关索引中的列表，您也可以轻松地做到这一点。