PROG-1106: นับเลข (Count) (Segment Tree)

#include <bits/stdc++.h>

using namespace std;
const int N = 2e5 + 10;
int tree[1 << 20];
bool del[N];
int n, k;

int Build(int idx, int l, int r){
    if(l == r) return tree[idx] = 1;
    int mid = (l + r) / 2;
    return tree[idx] = Build(2*idx, l, mid) + Build(2*idx + 1, mid + 1, r);
}

int Sum(int idx, int l, int r, int s, int e){
    if(l > e or r < s) return 0;
    if(s <= l and r <= e) return tree[idx];
    int mid = (l + r) / 2;
    return Sum(2*idx, l, mid, s, e) + Sum(2*idx + 1, mid + 1, r, s, e);
}

int Update(int idx, int l, int r, int pst){
    if(l == r) return tree[idx] = 0;
    int mid = (l + r) / 2;
    if(pst <= mid) return tree[idx] = Update(2*idx, l, mid, pst) + tree[2*idx + 1];
    else return tree[idx] = Update(2*idx + 1, mid + 1, r, pst) + tree[2*idx];
}

int Find(int s, int e, int x){
    int mn = e + 1, l = s, r = e;
    while(l <= r){
        int mid = (l + r) / 2;
        int sum = Sum(1, 1 , n, s, mid);
        if(Sum(1, 1, n, s, mid) >= x){
            mn = min(mn, mid);
            r = mid - 1;
        }
        else
            l = mid + 1;
    }
    return mn;
}

int main(){

    scanf("%d%d", &n, &k);

    Build(1, 1, n);

    int cur = 0, srch, sum;
    for(int sz=n;sz>1;sz--){
        /// search
        if(cur != 0) k = cur;
        if(k % sz == 0) srch = sz;
        else srch = k % sz;

        /// sum
        sum = 0;
        if(cur != n) sum = Sum(1, 1, n, cur + 1, n);

        /// find index
        if(sum >= srch) cur = Find(cur + 1, n, srch);
        else cur = Find(1, cur - 1, srch - sum);

        /// update index
        del[cur] = true;
        Update(1, 1, n, cur);
    }

    for(int i=1;i<=n;i++){
        if(!del[i]){
            printf("%d", i);
            return 0;
        }
    }

    return 0;
}