如何在Python中高效获取已排序数组元素的区间索引

问题背景 #

我们需要为已排序数组中的每个元素，找到其所属的区间索引（区间无重叠且整体有序，即np.c_[starts, ends].flatten()结果为已排序数组）。现有基于Numpy的实现通过两次searchsorted完成，但可利用输入数组的有序性进一步优化效率。

现有实现回顾 #

import numpy as np

def range_index(arr, starts, ends):
    _idx_starts = np.searchsorted(starts, arr, side='right')
    _idx_ends = np.searchsorted(ends, arr, side='left') + 1
    _idx_starts[_idx_starts != _idx_ends] = -1
    return _idx_starts

参数说明 #

• arr：待索引的已排序数组
• starts：已排序的区间起始值数组
• ends：已排序的区间结束值数组
• 区间特性：无重叠且整体有序（前一个区间的结束值 ≤ 后一个区间的起始值）

示例验证 #

_d = np.linspace(0, 20, 11)
_ranges = np.asarray([(1, 5), (7, 9), (12, 16)])
_idx = range_index(_d, _ranges[:, 0], _ranges[:, 1])
print(np.c_[_idx, _d])

输出：

[[-1.  0.]
 [ 1.  2.]
 [ 1.  4.]
 [-1.  6.]
 [ 2.  8.]
 [-1. 10.]
 [ 3. 12.]
 [ 3. 14.]
 [ 3. 16.]
 [-1. 18.]
 [-1. 20.]]

优化方案 #

方案1：单次二分搜索（减少一半二分开销） #

利用区间整体有序的特性，将所有区间的起始、结束值合并为一个有序边界数组，仅需一次searchsorted即可推导元素所属区间：

def optimized_range_index(arr, starts, ends):
    # 构造有序边界数组（符合输入条件，此数组天然有序）
    boundaries = np.c_[starts, ends].flatten()
    # 单次二分搜索获取插入位置
    pos = np.searchsorted(boundaries, arr, side='left')
    # 计算候选区间索引（从1开始计数）
    idx_candidate = (pos + 1) // 2
    # 标记不在任何区间的元素：插入位置为偶数时，元素落在区间间隙或两端外
    mask = (pos % 2 == 0)
    idx_candidate[mask] = -1
    # 返回整数类型索引，避免不必要的浮点转换
    return idx_candidate.astype(int)

原理说明

• 边界数组结构为[start0, end0, start1, end1, ..., startN-1, endN-1]
• 插入位置pos为奇数时：元素落在[start_i, end_i]区间内，对应索引(pos+1)//2
• 插入位置pos为偶数时：元素落在区间间隙或数组两端，标记为-1

方案2：线性扫描（O(n+m)时间复杂度） #

由于arr和区间均为有序，可使用双指针线性遍历，时间复杂度优于二分搜索的O(n log m)，适合大数组场景：

def linear_range_index(arr, starts, ends):
    # 初始化结果数组，默认值-1
    idx = np.full(len(arr), -1, dtype=int)
    m = len(starts)
    ptr = 0  # 区间遍历指针

    for i, x in enumerate(arr):
        # 跳过所有起始值大于当前元素的区间
        while ptr < m and starts[ptr] > x:
            ptr += 1
        # 指针越界，后续元素均无匹配区间，直接退出
        if ptr >= m:
            break
        # 检查当前元素是否在当前区间内
        if x <= ends[ptr]:
            idx[i] = ptr + 1
        else:
            # 元素大于当前区间结束值，移动指针寻找下一个可能的区间
            while ptr < m and ends[ptr] < x:
                ptr += 1
                if ptr < m and starts[ptr] <= x <= ends[ptr]:
                    idx[i] = ptr + 1
    return idx

原理说明

• 利用arr的有序性，区间指针仅需单向移动，无需回溯
• 遍历过程中直接匹配元素所属区间，避免重复二分计算

适用场景对比 #

• 单次二分搜索：适合区间数量较多（m较大）、数组长度中等的场景，代码简洁且内存占用低
• 线性扫描：适合数组长度极大（n远大于m）的场景，时间效率更高

验证结果 #

两种优化方案的输出与原实现完全一致：

_d = np.linspace(0, 20, 11)
_ranges = np.asarray([(1, 5), (7, 9), (12, 16)])

# 单次二分搜索验证
idx1 = optimized_range_index(_d, _ranges[:, 0], _ranges[:, 1])
print(np.c_[idx1, _d])

# 线性扫描验证
idx2 = linear_range_index(_d, _ranges[:, 0], _ranges[:, 1])
print(np.c_[idx2, _d])

输出均与原示例一致。

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