Question

我试图在numpy中对两个大型四维数组进行排序。

我想基于第一个数组的值2排序，并使用相同的索引对第二个数组进行排序。对于两个阵列，所有其他轴应保持相同的顺序。

以下代码执行我想要的操作，但依赖于在python中循环，所以它很慢。这些数组非常大，因此出于性能原因，我非常希望使用编译的numpy操作。或者其他一些方法可以编译这段代码（Cython？）。

import numpy as np

data = np.random.rand(10,6,4,1)
data2 = np.random.rand(10,6,4,3)

print data[0,0,:,:]
print data2[0,0,:,:]

for n in range(data.shape[0]):
  for m in range(data.shape[1]):

    sort_ids = np.argsort(data[n,m,:,0])

    data[n,m,:,:] = data[n,m,sort_ids,:]
    data2[n,m,:,:] = data2[n,m,sort_ids,:]


print data[0,0,:,:]
print data2[0,0,:,:]

Answer 1

也许有更好的解决方案，但这应该有效：

sort_ids = np.argsort(data,axis=2)

s1 = data.shape
s2 = data2.shape
d1 = data[np.arange(s1[0])[:,None,None,None],np.arange(s1[1])[None,:,None,None],sort_ids,np.arange(s1[3])[None,None,None,:]]
d2 = data2[np.arange(s2[0])[:,None,None,None],np.arange(s2[1])[None,:,None,None],sort_ids,np.arange(s2[3])[None,None,None,:]]

至少输出与您的代码相同。

Answer 2

找到一种方法来完成这项工作。它需要存储索引数组，这可能会给我带来一些内存问题，但速度要快一些。带时序比较的示例代码：

import numpy as np
import time

loops = 1000

data = np.random.rand(100,6,4,1)
data2 = np.random.rand(100,6,4,3)

start = time.time()
for n in range(loops):


  idxs = np.indices(data.shape)
  idxs2 = np.indices(data2.shape)

  sort_ids = np.argsort(data, 2)

  sorted_data = data[idxs[0], idxs[1], sort_ids, idxs[3]]
  sorted_data2 = data2[idxs2[0], idxs2[1], np.repeat(sort_ids, data2.shape[3], 3), idxs2[3]]

print 'Time Elapsed: %5.2f seconds' % (time.time() - start)



start = time.time()
for n in range(loops):

  sorted_data = np.zeros(data.shape)
  sorted_data2 = np.zeros(data2.shape)

  for n in range(data.shape[0]):
    for m in range(data.shape[1]):

      sort_ids = np.argsort(data[n,m,:,0])

      data[n,m,:,:] = data[n,m,sort_ids,:]
      data2[n,m,:,:] = data2[n,m,sort_ids,:]


print 'Time Elapsed: %5.2f seconds' % (time.time() - start)

Numpy - 按第一个数组的单轴排序两个ndarray

2 个答案: