对每个元素都是元组的 numpy 数组进行排序的最快方法是什么

Question

假设我有一个形状为 [N, 2] 的 numpy 数组。我想通过将每一行视为 2 个元素的元组来对数组进行排序。我只需要排序后的索引，而不需要返回排序后的数组。

由于 numpy 不直接支持此操作，因此我尝试通过以下方式使用结构化数组：

a = np.empty(N, dtype=[('f0', int), ('f1', int)])
a['f0'] = x
a['f1'] = y
index = np.argsort(a)

但是，我发现代码非常慢，比对标准 numpy 数组进行排序至少慢 100 倍。我很困惑为什么 argsort 操作对于结构化数组来说非常慢。我希望看到更快的实施。谢谢！

Answer 1

在 NumPy 中使用结构化数组对于某些操作（包括排序）可能效率低下，因为访问和管理结构化数据类型会产生相关开销。当您使用结构化数组时，NumPy 必须处理数据结构的额外复杂性，这可能会导致性能显着下降，尤其是对于大型数据集。

相反，您可以通过利用 NumPy 对常规数组进行排序的功能来实现所需的排序。这是一种更快、更有效的方法来获取原始二维数组的排序索引，而无需创建结构化数组。

您可以通过以下步骤实现此目的：

使用 np.lexsort 基于多个键执行稳定排序。 lexsort 函数首先根据最后一个键进行排序，依此类推。因此，您可以通过将每一行视为 2 个元素的元组来对 2D 数组进行排序。 以下是实现此方法的方法：

import numpy as np

# Sample data
N = 10  # Number of rows
x = np.random.randint(0, 100, size=N)  # First column
y = np.random.randint(0, 100, size=N)  # Second column

# Combine x and y into a 2D numpy array
a = np.column_stack((x, y))

# Get the sorted indices using lexsort
# Sort by the second column first, then by the first column
sorted_index = np.lexsort((a[:, 1], a[:, 0]))

# Display the original array and the sorted indices
print("Original array:")
print(a)
print("\nSorted indices:")
print(sorted_index)
print("\nSorted array (using sorted indices):")
print(a[sorted_index])

说明：

1. 数据生成：我们生成随机整数用于演示。在这里，我们使用 np.column_stack() 构造一个二维数组 a 来组合 x 和 y。

2. 排序索引：

np.lexsort() 用于按多个键排序。第一个参数是要排序的最后一个键，在本例中是第二列 (a[:, 1])。第二个参数是排序依据的第一个键，即第一列 (a[:, 0])。 这有效地将数组的行排序为两个元素的元组。 3.输出：

打印原始数组，以及排序后的索引和排序后的数组本身以进行验证。

此方法比使用结构化数组要快得多，并且是 NumPy 中按多列排序的推荐方法。