百万级一维NumPy数组：为每个元素寻找满足双条件的最小索引元素

问题明确 #

对arr1的每个元素arr1[i]，在arr2中找到最小的索引j>i，使得arr2[j]>arr1[i]，最终返回对应的arr2[j]（values数组）和j（indices数组）。

百万级数据下，广播、双重循环等O(n²)方法完全不可行，必须采用O(n log m)的高效算法，这里推荐线段树解法：

解法思路 #

1. 构建线段树：以arr2为基础，构建一棵存储区间最大值的线段树，快速判断任意区间内是否存在大于目标值的元素，避免无效遍历。
2. 逐个查询：对arr1的每个索引i，通过线段树查询arr2中i之后（j>i）第一个大于arr1[i]的元素索引j。
3. 收集结果：根据查询到的j，提取对应值和索引，生成结果数组。

代码实现 #

import numpy as np

class SegmentTree:
    def __init__(self, data):
        self.data = data
        self.n = len(data)
        self.size = 1
        # 找到大于等于n的最小2的幂次，用于线段树结构对齐
        while self.size < self.n:
            self.size <<= 1
        # 初始化最大值线段树，用NumPy提升效率
        self.max_tree = np.zeros(2 * self.size, dtype=data.dtype)
        # 填充叶子节点
        self.max_tree[self.size : self.size + self.n] = data
        # 构建上层节点，每个节点存储左右子节点的最大值
        for i in range(self.size - 1, 0, -1):
            self.max_tree[i] = max(self.max_tree[2*i], self.max_tree[2*i+1])
    
    def query_min_j(self, i, x):
        """查询最小的j > i，使得data[j] > x，不存在则返回-1"""
        return self._query(1, 0, self.size - 1, i, x)
    
    def _query(self, node, node_l, node_r, i, x):
        # 当前区间全部在i左侧，无符合条件的j
        if node_r <= i:
            return -1
        # 当前区间最大值<=x，无符合条件的元素
        if self.max_tree[node] <= x:
            return -1
        # 叶子节点，检查是否满足j>i
        if node_l == node_r:
            return node_l if node_l > i else -1
        mid = (node_l + node_r) // 2
        # 优先查询左子区间（保证找到最小的j）
        left_res = self._query(2*node, node_l, mid, i, x)
        if left_res != -1:
            return left_res
        # 左子区间无结果，查询右子区间
        return self._query(2*node+1, mid+1, node_r, i, x)

# 测试示例
arr1 = np.array([3,2,1,0,3,2,3])
arr2 = np.array([0,2,1,2,3,4,6,5])

st = SegmentTree(arr2)
values = []
indices = []
for i in range(len(arr1)):
    x = arr1[i]
    j = st.query_min_j(i, x)
    values.append(arr2[j])
    indices.append(j)

print(f"values = {values}")
print(f"indices = {indices}")

效率说明 #

• 构建时间：O(m)（m为arr2长度），百万级数据仅需数毫秒。
• 查询时间：单次查询O(log m)，百万次查询总操作约200万次，Python可轻松处理。
• 空间复杂度：O(m)，线段树需2倍于arr2长度的空间存储最大值，百万级数据仅占约8MB（float64类型），完全可控。

边界处理 #

若存在i使得arr2中所有j>i的元素都≤arr1[i]，query_min_j会返回-1，可根据需求填充默认值（如None或特定标记）。

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