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

#include <bits/stdc++.h>

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

void Update(int pst, int val){
    for(int i = pst; i <= n; i += i & -i)
        tree[i] += val;
}

int Sum(int pst, int sum = 0){
    for(int i = pst; i > 0; i -= i & -i)
        sum += tree[i];
    return sum;
}

int FindSum(int s, int e){
    int sums = Sum(s - 1);
    int sume = Sum(e);
    return sume - sums;
}

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

int main(){

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

    for(int i=1;i<=n;i++)
        Update(i, 1);

    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 = FindSum(cur + 1, n);

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

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

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

    return 0;
}