如何高效获取二维查询数组在已排序二维Numpy数组的各行插入位置

背景与单查询处理方式 #

现有一个沿轴0排序的二维Numpy数组R（50k行、1k列），对于单个一维查询数组x，可通过以下代码获取x中每个值在R各列的插入位置：

import numpy as np

R = np.array([[0, 8, 12, 2],
              [3, 16, 13, 5],
              [4, 19, 13, 11],
              [6, 23, 16, 12]])

x = np.array([4, 9, 17, 6])

# 期望输出: array([2, 1, 4, 2])
idx = (R < x).sum(axis=0)

需求与现有方案的问题 #

当前需要处理二维查询数组X（每行是一个查询数组，行数可与R相当），此前尝试的两种方法存在明显缺陷：

R = np.array([[0, 8, 12, 2],
              [3, 16, 13, 5],
              [4, 19, 13, 11],
              [6, 23, 16, 12]])

X = np.array([[4, 9, 17, 6],
              [1, 2, 12, 20]])

# 期望输出:
# array([[2, 1, 4, 2],
#        [1, 0, 0, 4]])

# 方案1：内存占用过高
# 会生成形状为(M, 50k, 1k)的布尔数组，内存开销巨大
idx = (R < X[:, None]).sum(axis=1)

# 方案2：速度过慢
# Python循环遍历X的每一行，无法利用Numpy的向量化优化
idx = np.array([(R < x).sum(axis=0) for x in X])

尝试过searchsorted2d但因循环调用速度仍不理想，需要更高效的向量化通用方案。

优化方案：利用np.searchsorted的广播与向量化特性 #

由于R沿轴0排序，其每一列都是升序数组。我们可以通过转置R，结合np.searchsorted的axis参数实现全向量化计算：

import numpy as np

R = np.array([[0, 8, 12, 2],
              [3, 16, 13, 5],
              [4, 19, 13, 11],
              [6, 23, 16, 12]])
X = np.array([[4, 9, 17, 6],
              [1, 2, 12, 20]])

# 转置R，将原数组的列转换为行，形状变为(1k, 50k)
R_transposed = R.T

# 使用searchsorted，指定axis=1，在R_transposed的每一行（原R的列）中查找X对应位置值的插入点
idx = np.searchsorted(R_transposed, X, axis=1)

print(idx)
# 输出:
# [[2 1 4 2]
#  [1 0 0 4]]

方案优势 #

• 内存高效：避免了生成庞大的中间布尔数组，仅利用R的转置和广播机制，内存占用远低于方案1。
• 速度极快：np.searchsorted基于底层优化的C实现，无Python循环开销，比方案2的循环方式效率提升显著。
• 通用性强：只要R的列是排序状态（题目已保证），该方法适用于任意规模的R和X，无需额外修改。

内容的提问来源于stack exchange，提问作者Sakhawat Hossain