Search code examples
c++algorithmdata-structuressegment-tree

How to get the first value greater than x in a range with segment tree


I am trying to use a segment tree to handle queries (in O(log N) per query) of the form: find the first value > x in a particular range [ql, qr] (of an array of size N).

I've been using the recursive method to create the segment tree and to update it, but I wasn't able to retrieve the index of the first value > x in my range.

Here is my implementation that doesn't currently work:

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

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

//recursive build of the segtree
long long int build(int i, int l, int r, vector<int> &torri) {

    if(r-l==1) return seg[i] = torri[l]; //base case

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

    //return the value of the segment
    return seg[i];
}

//recursive update of the segtree
void update(int i, int l, int r, int pos, long long int val) {

    if(r-l==1) { seg[i] = val; return; } //base case

    int mid = (l+r)/2; //find the middle
    if(pos<mid) update(2*i, l, mid, pos, val); //left half
    else update(2*i+1, mid, r, pos, val); //right half

    seg[i] = max(seg[2*i],seg[2*i+1]); //update with the max
}

int query_destra (int i, int l, int r, int ql, int qr, int x) {

    //base case
    if (ql >= r or qr <= l) return -1;//no intersection

    if (l >= ql and r <= qr) { //included
        if(seg[i] <= x) return -1; //if ... the node is not interesting

        //return i; <----------- IM MISSING SOMETHING HERE TO MAKE IT RETURN I CORRECTLY
    }

    int mid = (l+r)/2; //find the center

    //prioritizes the first value, if its not interesting, go to the second
    int ans = query_destra(2*i, l, mid, ql, qr, x);
    if(ans != -1) return ans;
    return query_destra(2*i+1, mid, r, ql, qr, x);
}

pair<int, int> chiedi(int x) {
    return make_pair(0 /*query_sinistra(1,0,sizeN, 0, x, torre[x])*/,query_destra(1, 0, sizeN, x, sizeN, torre[x]));
}

void cambia(int x, int h) {
    update(1, 0, sizeN, x, h);
}

void inizializza(int N, vector<int> H) {
    sizeN = H.size();
    torre.resize(sizeN);
    seg.resize(sizeN*4);
    for(int i=0; i<sizeN; i++) torre[i] = H[i];
    build(1,0,N,H);
}

int main() {
    vector vec{6, 4, 2, 0, 8, 4};
    inizializza(vec.size(), vec);
    cout << query_destra(1, 0, sizeN, 1, 5, 4) << '\n'; 
    // should produce 4, the index of the first value larger than 4 in the subarray from index 1 to 5 (exclusive), but gives -1
}

I've been able to retrieve the highest in the range, but not the first, so I tried a different approach: I tried simply putting "return i" inside the included cases but from the results it gave me it was VERY wrong, and not putting somewhere "return i" simply makes it return -1 each time.

After failing on my own, I tried following what other people did, but not only I wasn't able to understand most of the code they used, but it also didn't work for what I needed to do.

An example of how the result should be, with a starting array of [6, 4, 2, 0, 8, 4], with x = 4 and my range going from position 1 to position 5, my program should be able to return 4 (the position of the 8, the first number higher than the starting one). I have to do this thing for a number in the array both forwards and backwards. I'm using C++17. Any help is appreciated, thanks in advance.


Solution

  • Your solution is almost correct; you are just missing the point at which you should return the actual index. When you have narrowed your search to a leaf node while descending the tree, you can directly return the left index.

    Replace these lines

    if (l >= ql and r <= qr) {
        if(seg[i] <= x) return -1;
    
        //return i;
    }
    

    with

    if (seg[i] <= x) return -1;
    if (r == l + 1) return l;