Codeforces 1181D Irrigation

Irrigation

把询问离线， 从小到大解决， 转换成求第k大的问题， 套个平衡树就好啦。

#include<bits/stdc++.h>
#include <bits/extc++.h>
#define LL long long
#define LD long double
#define ull unsigned long long
#define fi first
#define se second
#define mk make_pair
#define PLL pair<LL, LL>
#define PLI pair<LL, int>
#define PII pair<int, int>
#define SZ(x) ((int)x.size())
#define ALL(x) (x).begin(), (x).end()
#define fio ios::sync_with_stdio(false); cin.tie(0);

using namespace std;
using namespace __gnu_pbds;

const int N = 5e5 + 7;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const int mod = 998244353;
const double eps = 1e-8;
const double PI = acos(-1);

template<class T, class S> inline void add(T& a, S b) {a += b; if(a >= mod) a -= mod;}
template<class T, class S> inline void sub(T& a, S b) {a -= b; if(a < 0) a += mod;}
template<class T, class S> inline bool chkmax(T& a, S b) {return a < b ? a = b, true : false;}
template<class T, class S> inline bool chkmin(T& a, S b) {return a > b ? a = b, true : false;}

template <class T>
using Tree = tree<T, null_type, std::less<T>, rb_tree_tag,tree_order_statistics_node_update>;
Tree<int> bst;

int n, m, q;
PLL qus[N];
PLL now[N];

int a[N];
int cnt[N];
int ans[N];

int main() {
    scanf("%d%d%d", &n, &m, &q);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &a[i]);
        now[a[i]].fi++;
    }
    for(int i = 1; i <= m; i++) now[i].se = i;
    sort(now + 1, now + 1 + m);

    for(int i = 1; i <= q; i++) {
        scanf("%lld", &qus[i].fi);
        qus[i].se = i;
    }
    sort(qus + 1, qus + 1 + q);

    LL pre = 0;
    LL cnt = 0;
    LL tmp = n;

    for(int i = 1, j = 1; i <= q; i++) {
        LL k = qus[i].fi;
        while(j <= m && (now[j].fi - pre) * cnt + tmp < k) {
            tmp += (now[j].fi - pre) * cnt;
            cnt++;
            bst.insert(now[j].se);
            pre = now[j].fi;
            j++;
        }
        LL need = k - tmp;
        LL pos = need % cnt;

        if(pos) pos = *bst.find_by_order(pos - 1);
        else pos = *bst.find_by_order(cnt - 1);

        ans[qus[i].se] = pos;
    }
    for(int i = 1; i <= q; i++) printf("%d
", ans[i]);
    return 0;
}

/*
*/