Submission #738480

#	Time	Username	Problem	Language	Result	Execution time	Memory
738480		lorenzoferrari	Regions (IOI09_regions)	C++17	40 / 100	8032 ms	34024 KiB

regions

#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#include "bits/stdc++.h"
using namespace std;
using LL = long long;

static constexpr int SQ = 500;

struct Segtree {
    int n;
    vector<int> t;
    Segtree(int n) : n(n), t(2*n) {}
    void update(int p, int v) {
        for (t[p += n] += v; p > 1; p >>= 1) {
            t[p >> 1] = t[p] + t[p ^ 1];
        }
    }
    int query(int l, int r) {
        int ans = 0;
        for (l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if (l & 1) ans += t[l++];
            if (r & 1) ans += t[--r];
        }
        return ans;
    }
};

int main() {
    int n; cin >> n;
    int r; cin >> r;
    int q; cin >> q;
    vector<int> h(n), p(n), f(r);
    vector<vector<int>> adj(n);
    cin >> h[0]; ++f[--h[0]];
    for (int i = 1; i < n; ++i) {
        cin >> p[i] >> h[i];
        --p[i], --h[i];
        adj[p[i]].push_back(i);
        ++f[h[i]];
    }
    vector<int> ff;
    vector<int> fi(r, -1);
    for (int i = 0; i < r; ++i) {
        if (f[i] >= SQ) {
            fi[i] = ff.size();
            ff.push_back(i);
        }
    }
    int nn = ff.size();

    vector<int> in(n), out(n), ord;
    { // linearize the tree
        int t = 0;
        auto dfs = [&](auto&& self, int v) -> void {
            in[v] = t++;
            ord.push_back(v);
            for (int u : adj[v]) {
                self(self, u);
            }
            out[v] = t;
        };
        dfs(dfs, 0);
    }

    vector<vector<int>> vs(r);
    {
        for (int v : ord) {
            int c = h[v];
            vs[h[v]].push_back(v);
        }
    }

    vector<vector<LL>> sub(r, vector<LL>(nn));
    vector<vector<LL>> anc(r, vector<LL>(nn));
    {
        auto dfs = [&](auto&& self, int v, int i, int acc = 0) -> int {
            int cnt = 0;
            for (int u : adj[v]) {
                cnt += self(self, u, i, acc + (h[v] == ff[i]));
            }
            sub[h[v]][i] += cnt;
            anc[h[v]][i] += acc;
            return cnt + (h[v] == ff[i]);
        };
        for (int i = 0; i < nn; ++i) {
            dfs(dfs, 0, i);
        }
    }

    Segtree st(n);
    for (int r1, r2; q--;) {
        cin >> r1 >> r2;
        --r1, --r2;
        if (f[r2] >= SQ) {
            cout << sub[r1][fi[r2]] << endl;
        } else if (f[r1] >= SQ) {
            cout << anc[r2][fi[r2]] << endl;
        } else {
            LL ans = 0;
            for (int v : vs[r2]) {
                st.update(in[v], 1);
            }
            for (int u : vs[r1]) {
                ans += st.query(in[u]+1, out[u]);
            }
            for (int v : vs[r2]) {
                st.update(in[v], -1);
            }
            cout << ans << endl;
        }
    }
}

Compilation message (stderr)

regions.cpp: In function 'int main()':
regions.cpp:68:17: warning: unused variable 'c' [-Wunused-variable]
   68 |             int c = h[v];
      |                 ^

• Subtask 1: 15 / 15
• Subtask 2: 25 / 85