如何使用 np.argsort 的二维数组对 lexsort 生成相同的结果？

Question

我已经使用可以与 lexsort 匹配的一维数组 argsort 进行了复制。

#a = 1d np.array
#b = 1d np.array

def lexsort_copy(a,b):
    idxs= np.argsort(a,kind='stable')
    return idxs[np.argsort(b[idxs],kind='stable')]

lexsort_copy(a,b) == np.lexsort((a,b))

这给了我相同的输出，但我正在努力如何使用二维数组复制它。

测试二维数组：

test=np.array([[100,100,100,100,111,400,120],[229,1133,152,210,120,320,320]])
np.lexsort(test)

输出：

array([4, 2, 3, 0, 6, 5, 1], dtype=int64)

我们如何在不对二维数组使用 lexsort 的情况下复制上面的输出？

如有任何解决方案，我们将不胜感激！谢谢！

Answer 1

您可以扩展 lexsort_copy 以使用二维数组，如下所示：

def lexsort2D_copy(data):
  idxs = np.arange(data.shape[1])
  for i in range(data.shape[0]-1):
    idxs = np.argsort(data[i][idxs],kind='stable')
  return idxs[np.argsort(data[-1][idxs],kind='stable')]

测试：

test=np.array([[100,100,100,100,111,400,120],
               [229,1133,152,210,120,320,320],
               [29,133,12,10,10,20,3120]])
np.lexsort(test)  == lexsort2D_copy(test)

test=np.array([[100,100,100,100,111,400,120],
               [229,1133,152,210,120,320,320]])
np.lexsort(test)  == lexsort2D_copy(test)

输出：

array([ True,  True,  True,  True,  True,  True,  True])
array([ True,  True,  True,  True,  True,  True,  True])

如何使用 np.argsort 的二维数组对 lexsort 生成相同的结果？

how to generate same result using 2D array of np.argsort to lexsort?

python

arrays

sorting

numpy