使用自定义谓词排序numpy数组

Question

我想使用在第二维向量(大小:4)上运行的自定义谓词,沿着第一维(大小:n)对我的numpy形状[n,4]数组进行排序.我想做的C++版本如下,真的非常简单.我已经看到如何使用python列表执行此操作,但我找不到使用numpy数组执行此操作的语法.这可能吗？关于np的文档.排序,np.argsort,np.lexsort没有提到自定义谓词.

// c++ version
vector< float[4] > v = init_v(); 
float[4] p = init_p();
std::sort(v.begin(), v.end(), [&p](const auto& lhs, const auto& rhs) {
   return myfn(p, lhs) > myfn(p, rhs); });

编辑:下面是我想用于排序的python代码.即对于我的阵列的每个'行'(n:4),我将计算欧几里德3D距离(即仅前3列)的平方到固定点.

# these both operate on numpy vectors of shape [4] (i.e. a single row of my data matrix)
def dist_sq(a,b):
    d = a[:3]-b[:3]
    return np.dot(d*d)

def sort_pred(lhs, rhs, p):
    return dist_sq(lhs, p) > dist_sq(rhs, p)

Answer 1

在numpy中,您将(向量化)顺序定义函数应用于数组,然后用于np.argsort按结果排序.

这比C++版本的空间效率低,但这就是你通常用numpy实现性能的方式.

import numpy as np    

def myfn(x):
    return np.sin(x[:, 1])  # example: sort by the sine of the second column

a = np.random.randn(10, 4)

predicate = myfn(a)  # not sure if predicate is the best name for this variable
order = np.argsort(predicate)

a_sorted = a[order]