[英]Numpy - Sorting two ndarrays by single axis of first array
我正在嘗試在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,:,:]
也許有更好的解決方案,但這應該可行:
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,:]]
至少輸出與您的代碼相同。
找到了一種使這項工作有效的方法。 它需要存儲一個索引數組,這可能會給我帶來一些內存問題,但是速度更快。 帶有時間比較的示例代碼:
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)
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.