如何对numpy结构化数组执行<等有序比较操作？

问题背景 #

用元组dtype定义的Numpy结构化数组，比如：

import numpy as np

A = np.array([(3, 2),
              (1, 8),
              (5, 1),
              (3, 1),
              (4, 7)],
             dtype='i8,i8')

支持np.sort(A)得到字典序结果，也能做向量化的==/!=比较，但直接执行A < (3, 2)这类有序比较会触发TypeError。手动逐列实现字典序判断扩展性差，而对超大数组全量排序再取前N个，耗时过高。

解决方案 #

1. 用np.lexsort生成索引，按需提取

np.lexsort会基于多列生成字典序的索引，仅排序索引而非整个数组，效率优于全量排序。示例提取前2个字典序最小元素：

# 获取结构化数组的所有字段列，注意lexsort按最后一个传入的键优先排序，需反转列顺序
keys = [A[field] for field in A.dtype.names]
sorted_idx = np.lexsort(keys[::-1])

# 提取前2个元素
top_elements = A[sorted_idx[:2]]

输出：

array([(1, 8), (3, 1)], dtype=[('f0', '<i8'), ('f1', '<i8')])

2. 结合np.argpartition实现O(n)级部分排序

如果仅需前K个最小元素，np.argpartition的时间复杂度为O(n)，远快于全排序的O(n log n)。可以构造复合键来映射字典序：

# 构造复合键：第一列权重远大于第二列，确保字典序优先级
max_val = np.iinfo(np.int64).max // 2  # 避免数值溢出
composite_key = A['f0'] * max_val + A['f1']

# 找到前2个最小复合键的索引
partition_idx = np.argpartition(composite_key, 2)[:2]
# 对提取出的元素做小范围排序，保证结果是严格字典序
top_elements = np.sort(A[partition_idx])

3. 借助Pandas简化实现（若允许引入依赖）

Pandas对结构化数据的字典序支持更友好，且内部优化了排序性能：

import pandas as pd

df = pd.DataFrame(A)
top_elements = df.nsmallest(2, columns=df.columns).to_records(index=False)

代码简洁，超大数组场景下性能表现优异。

补充说明 #

Numpy结构化数组未实现默认的字典序有序比较（<, <=等），因为字段可能包含不同数据类型，无法统一处理；但np.sort内部使用了lexsort逻辑，因此能正确返回字典序结果。

内容的提问来源于stack exchange，提问作者Dan R