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

Question

我想对形状为[n,4]的numpy数组进行排序，沿第一维（大小：n）使用作用于第二维向量（大小：4）的自定义谓词。下面是我想要做的C++版本，实际上非常简单。我已经看到了如何使用python lists来完成这个任务，但我找不到使用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代码。即对于我的数组的每一行（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)

- memo

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]

- MB-F

---

这是一个不错的方法，但它不能处理按多个键排序的情况 --- 即首先按列A排序，然后按列B等次要排序。 - Kyle Pena

@KylePena 你可以使用numpy.lexsort代替numpy.argsort。直接传递数组或者像这个问题/答案中那样应用谓词转换即可。 - MB-F

这并没有解决问题。@memo想要根据比较两个对象的函数来排序一组对象。Argsort仅适用于具有某些数字表示形式的对象。 - Josh Albert

@JoshAlbert 这个回答在问题被编辑之前就已经回答了。而且，显然这个答案对于提问者来说是有效的。如果您认为这个答案不够充分，可以自由地发布另一个答案。 - MB-F