使用自定义谓词对numpy数组进行排序 [英] sort numpy array with custom predicate

问题描述

我想使用对第二维矢量(size:4)进行操作的自定义谓词，沿第一维(size:n)对形状为[n，4]的numpy数组进行排序.下面是我想做的C ++版本，真的很简单.我已经看过如何使用 python列表，但是我找不到使用numpy数组执行该操作的语法.这可能吗?关于np的文档. sort ，np. argsort ，np.

I'd like to sort my numpy array of shape [n,4], along first dimension (size:n) using a custom predicate operating on the 2nd dimension vector (size:4). The C++ version of what I'd like to do is below, it's quite simple really. I've seen how to do this with python lists, but I can't find the syntax to do it with numpy arrays. Is this possible? The documentation on np.sort, np.argsort, np.lexsort doesn't mention custom predicates.

// 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列)到固定点的平方.

below is the python code I would like to use for the sorting. I.e. for each 'row' (n:4) of my array, I'd calculate the square of the euclidean 3D distance (i.e. only the first 3 columns) to a fixed point.

# 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)

推荐答案

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

In numpy you would apply the (vectorized) order defining function to the array, then use np.argsort to sort by the result.

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

This is less space efficient than the C++ version, but that is how you usually achieve performance with 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]

这篇关于使用自定义谓词对numpy数组进行排序的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！