如何用线段树查询区间内第一个大于x的元素索引

问题背景 #

需要用线段树实现数组指定区间[ql, qr]内第一个大于x的元素索引的查询，要求单次查询时间复杂度O(logN)。已实现递归式线段树的构建与更新，但查询逻辑错误，无法返回正确索引。例如对数组[6, 4, 2, 0, 8, 4]查询区间[1,5)中第一个大于4的元素，预期返回索引4（对应元素8），但当前代码返回-1。同时需要支持**正向（从左到右找第一个）和反向（从右到左找第一个）**两种查询。

错误核心原因 #

当前query_destra函数的逻辑缺陷：

• 当线段树节点完全落在查询区间内且最大值大于x时，没有继续递归到叶子节点定位具体数组索引，直接跳过该分支，导致后续递归无法找到目标元素。
• 错误地尝试返回线段树节点的索引i，但线段树节点索引与原数组元素索引不对应，节点i代表的是一段区间，不是单个元素的位置。

修正后的完整代码 #

#include <utility>
#include <iostream>
#include <vector>
using namespace std;

vector<int> seg;
vector<int> torre;
int sizeN;

// 递归构建线段树，存储区间最大值
long long int build(int i, int l, int r, vector<int> &torri) {
    if (r - l == 1) return seg[i] = torri[l]; // 叶子节点，对应原数组单个元素

    int mid = (l + r) / 2;
    seg[i] = max(build(2 * i, l, mid, torri), build(2 * i + 1, mid, r, torri));
    return seg[i];
}

// 递归更新线段树指定位置的值
void update(int i, int l, int r, int pos, long long int val) {
    if (r - l == 1) {
        seg[i] = val;
        return;
    }

    int mid = (l + r) / 2;
    if (pos < mid)
        update(2 * i, l, mid, pos, val);
    else
        update(2 * i + 1, mid, r, pos, val);

    seg[i] = max(seg[2 * i], seg[2 * i + 1]);
}

// 正向查询：在[ql, qr)中从左到右找第一个大于x的元素索引
int query_left(int i, int l, int r, int ql, int qr, int x) {
    // 无交集，返回-1
    if (ql >= r || qr <= l) return -1;

    // 当前节点完全在查询区间内
    if (l >= ql && r <= qr) {
        // 区间最大值不大于x，直接返回-1
        if (seg[i] <= x) return -1;
        // 叶子节点，返回原数组索引l
        if (r - l == 1) return l;
    }

    int mid = (l + r) / 2;
    // 优先查左子树，找到就返回（保证是第一个）
    int left_ans = query_left(2 * i, l, mid, ql, qr, x);
    if (left_ans != -1) return left_ans;
    // 左子树没找到，查右子树
    return query_left(2 * i + 1, mid, r, ql, qr, x);
}

// 反向查询：在[ql, qr)中从右到左找第一个大于x的元素索引
int query_right(int i, int l, int r, int ql, int qr, int x) {
    // 无交集，返回-1
    if (ql >= r || qr <= l) return -1;

    // 当前节点完全在查询区间内
    if (l >= ql && r <= qr) {
        if (seg[i] <= x) return -1;
        if (r - l == 1) return l;
    }

    int mid = (l + r) / 2;
    // 优先查右子树，找到就返回（保证是最右边的第一个）
    int right_ans = query_right(2 * i + 1, mid, r, ql, qr, x);
    if (right_ans != -1) return right_ans;
    // 右子树没找到，查左子树
    return query_right(2 * i, l, mid, ql, qr, x);
}

pair<int, int> chiedi(int x) {
    return make_pair(
        query_left(1, 0, sizeN, 0, x, torre[x]), // 正向查询左半区间[0,x)
        query_right(1, 0, sizeN, x, sizeN, torre[x]) // 反向查询右半区间[x,sizeN)
    );
}

void cambia(int x, int h) {
    update(1, 0, sizeN, x, h);
    torre[x] = h; // 同步更新原数组
}

void inizializza(int N, vector<int> H) {
    sizeN = H.size();
    torre = H; // 直接赋值，简化代码
    seg.resize(sizeN * 4);
    build(1, 0, N, H);
}

int main() {
    vector vec{6, 4, 2, 0, 8, 4};
    inizializza(vec.size(), vec);
    // 测试正向查询[1,5)中第一个大于4的元素，预期返回4
    cout << query_left(1, 0, sizeN, 1, 5, 4) << '\n';
    // 测试反向查询[1,5)中最后一个大于4的元素，预期返回4
    cout << query_right(1, 0, sizeN, 1, 5, 4) << '\n';
    // 修改索引4的值为3，再次查询[1,5)中第一个大于4的元素，预期返回-1
    cambia(4, 3);
    cout << query_left(1, 0, sizeN, 1, 5, 4) << '\n';
    return 0;
}

关键修正说明 #

1. 正向查询逻辑修正：
   • 当节点完全在查询区间内且最大值大于x时，先判断是否为叶子节点，是则直接返回原数组索引l；否则继续递归左子树，左子树找到结果就返回，确保是区间内最左边的符合条件的元素。
2. 新增反向查询函数：
   • 递归顺序调整为优先查询右子树，再查左子树，保证返回的是区间内最右边的第一个符合条件的元素。
3. 原数组同步更新：
   • 在cambia函数中同步更新torre数组，避免后续查询时使用旧值。
4. 代码简化：
   • 初始化时直接用torre = H替代循环赋值，更简洁。

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