numpy:将argsort应用于数组

时间:2012-06-28 21:55:41

标签: python arrays numpy

argsort()函数返回一个索引矩阵,可用于索引原始数组,以使结果与sort()结果匹配。

有没有办法应用这些指数?我有两个数组,一个是用于获取排序顺序的数组,另一个是一些关联数据。

我想计算assoc_data[array1.argsort()],但这似乎不起作用。

以下是一个例子:

z=array([1,2,3,4,5,6,7])
z2=array([z,z*z-7])
i=z2.argsort()

z2=array([[ 1,  2,  3,  4,  5,  6,  7],
          [-6, -3,  2,  9, 18, 29, 42]])
i =array([[1, 1, 1, 0, 0, 0, 0],
          [0, 0, 0, 1, 1, 1, 1]])

我想将i应用于z2(或其他具有关联数据的数组),但我不知道该怎么做。

5 个答案:

答案 0 :(得分:8)

这可能是矫枉过正,但这将适用于第二种情况:

import numpy as np
axis = 0
index = list(np.ix_(*[np.arange(i) for i in z2.shape]))
index[axis] = z2.argsort(axis)
z2[index]

# Or if you only need the 3d case you can use np.ogrid.

axis = 0
index = np.ogrid[:z2.shape[0], :z2.shape[1], :z2.shape[2]]
index[axis] = z2.argsort(axis)
z2[index]

答案 1 :(得分:4)

你很幸运我刚获得了numpyology的硕士学位。

>>> def apply_argsort(a, axis=-1):
...     i = list(np.ogrid[[slice(x) for x in a.shape]])
...     i[axis] = a.argsort(axis)
...     return a[i]
... 
>>> a = np.array([[1,2,3,4,5,6,7],[-6,-3,2,9,18,29,42]])
>>> apply_argsort(a,0)
array([[-6, -3,  2,  4,  5,  6,  7],
       [ 1,  2,  3,  9, 18, 29, 42]])

有关正在发生的事情的解释,请参阅我对this question的回答。

答案 2 :(得分:1)

使用np.take_along_axis

np.take_along_axis(z2, i, axis=1)
Out[31]: 
array([[ 1,  2,  3,  4,  5,  6,  7],
       [-6, -3,  2,  9, 18, 29, 42]])

答案 3 :(得分:0)

啊,阿哈,想通了。

In [274]: z2[i,range(z2.shape[1])]
Out[274]:
array([[-6, -3,  2,  4,  5,  6,  7],
       [ 1,  2,  3,  9, 18, 29, 42]])

答案 4 :(得分:0)

我最近遇到了这个问题(使用a.argsort(axis=something)b进行排序)。这是我的实现,如果无法负担运行时间,则没有列表:

#For the example
import numpy as np
shape=(2,2,3,2)
t=np.random.randint(10,size=np.array(shape).prod()).reshape(shape)

#The axis to sort by
sortby=-1

#The actual sorting
x=t.swapaxes(sortby,-1)  #I actually only know to do the final dimension, so I cheat
a=np.argsort(x,axis=-1)
p = np.array(x.shape[:-1]).prod()
x.T[a.T].T.reshape(p**2,x.shape[-1])[::p+1].reshape(x.shape).swapaxes(sortby,-1)

我不确定为什么这样做,但似乎适用于任何尺寸和排序轴。