python - 根据排列数组从一组数组中选择值

Question

我有 3 个形状为 2xN 的 numpy 数组(N 很大，有几百万)，称它们为 a1、a2、a3。然后我有另一个形状为 Nx3 的数组，其行值引用数组 a1、a2、a3 之一，称之为排列。这个排列数组看起来像:[[0, 1, 2],[1,2,0],[1,0,2],...最多 N 行]

我想创建另外 3 个形状为 2xN 的 numpy 数组 b1、b2、b3，它们具有原始 a1、a2、a3 的内容，但它们的列已根据排列数组的行进行了排列。

我尝试过花哨的索引堆叠 3 个数组和 numpy.choose，但我无法让它工作。我正在寻找没有 python 循环的解决方案。任何帮助将不胜感激!

编辑

为了澄清，我展示了我正在尝试做的事情的 python 循环实现:

aa = np.dstack((a1, a2, a3))
bb = np.empty_like(aa)
for i, o in enumerate(permutations):
    bb[:,i, np.arange(3)] = aa[:, i, o]

然后我会从 bb 中检索 b1、b2、b3。

Answer 1

与 fancy-indexing , 你可以做 -

bb = aa[:,np.arange(N),permutations.T]

请注意，这将是 (2,3,N) 的形状。因此，要选择 b1、b2、b3，您可以:

b1,b2,b3 = bb[:,0,:], bb[:,1,:], bb[:,2,:]

或者如果您坚持 bb 的形状与发布的代码相同，您可以添加:

bb = bb.swapaxes(1,2)

---

这是另一种方法，使用线性索引、切片，当然还有 NumPy broadcasting -

idx = permutations + 3*np.arange(N)[:,None]    
bb = aa.reshape(2,-1)[:,idx].reshape(2,N,3)

这将创建一个 bb，其形状与发布的循环代码相同。

---

运行时测试

In [189]: def original_app(aa,permutations):
     ...:     bb = np.empty_like(aa)
     ...:     for i, o in enumerate(permutations):
     ...:         bb[:,i, np.arange(3)] = aa[:, i, o]
     ...:     return bb
     ...: 
     ...: 
     ...: def linear_index_app(aa,permutations):
     ...:     idx = permutations + 3*np.arange(N)[:,None]    
     ...:     return aa.reshape(2,-1)[:,idx].reshape(2,N,3)
     ...: 

In [190]: # Setup input arrays
     ...: N = 10000
     ...: a1 = np.random.rand(2,N)
     ...: a2 = np.random.rand(2,N)
     ...: a3 = np.random.rand(2,N)
     ...: 
     ...: permutations = np.random.randint(0,3,(N,3))
     ...: aa = np.dstack((a1, a2, a3))


In [191]: %timeit original_app(aa,permutations)
10 loops, best of 3: 128 ms per loop

In [192]: %timeit aa[:,np.arange(N),permutations.T]
1000 loops, best of 3: 972 µs per loop

In [193]: %timeit linear_index_app(aa,permutations)
1000 loops, best of 3: 1.02 ms per loop

所以，看来 fancy-indexing 是其中最好的一个!