根据numpy数组中的条件创建新的numpy数组

Question

我想知道如何根据另一个数组中的特定条件来创建一个新数组。举个例子，如果我有一个数组，它包含一个基础数字、一个起始数字和一个结束数字，还有多个其他的基础数字。我想创建一个新的矩阵，里面列出基础数字、循环次数（根据起始和结束数字）以及与之相关的其他基础数字，同时忽略0。我的目标是找到一种方法来实现这一点，而不使用for循环。

比如，我想从数组A中得到数组B。

           Base Start End Base1 Base2 Base3
A=np.array([[100,  1,   2,  101,  102,  103],
            [101,  3,   4,  100,  103,    0]])


B=np.array[[100,1,101,1],
           [100,1,102,1],
           [100,1,103,1],
           [100,2,101,2],
           [100,2,102,2],
           [100,2,103,2],
           [101,3,100,3],
           [101,3,103,3],
           [101,4,100,4],
           [101,4,103,4]]

谢谢你的帮助！

Answer 1

给你看看：完全使用numpy，没有在Python层面使用for循环

C = np.empty((len(A), 2,3,4), np.int)
C[...,0] = A[:,0  ][:,None,None]
C[...,1] = A[:,1:3][:,:   ,None]
C[...,2] = A[:,3: ][:,None,:   ]
C[...,3] = np.arange(len(A)*2).reshape(len(A),2,1)+1

C = C.reshape(-1,4)
C = C[C[:,2]!=0]

print C

编辑：稍微整理了一下

编辑2：或者为了更简洁和神秘（见评论）：

C = np.empty((len(A), 2,3,2,2), np.int)
C[...,0,0] = A[:,0  ][:,None,None]
C[...,1,0] = A[:,3: ][:,None,:   ]
C[...,:,1] = A[:,1:3][:,:   ,None,None]

Answer 2

给你看看……这个列表推导式真是复杂。

A = array([[100,   1,   2, 101, 102, 103], [101,   3,   4, 100, 103,   0]])
B = [[A[i,0], b, c, b] for i in range(len(A)) for b in A[i,1:3] for c in A[i,3:6] if c != 0]

>>> B
[[100, 1, 101, 1],
 [100, 1, 102, 1],
 [100, 1, 103, 1],
 [100, 2, 101, 2],
 [100, 2, 102, 2],
 [100, 2, 103, 2],
 [101, 3, 100, 3],
 [101, 3, 103, 3],
 [101, 4, 100, 4],
 [101, 4, 103, 4]]