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

Question

我想使用对第二维向量（大小：4）进行操作的自定义谓词，沿第一维（大小：n）对形状 [n,4] 的 numpy 数组进行排序。下面是我想要做的 C++ 版本，它确实非常简单。我已经了解了如何使用 python 列表 执行此操作，但我找不到使用 numpy 数组执行此操作的语法。这可能吗？ np.sort、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代码。 IE。对于数组的每个“行”(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]