Submission #1145631

#	Time	Username	Problem	Language	Result	Execution time	Memory
1145631		monkey133	Team Coding (EGOI24_teamcoding)	C++20	0 / 100	0 ms	328 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
// Query structure for Mo's algorithm.
struct Query {
    int L, R, idx;
};
 
int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    // Input:
    // First line: n (number of nodes) and q (number of queries)
    int n, q;
    cin >> n >> q;
    
    // Read node colors.
    // (Assume colors are integers; adjust if needed.)
    vector<int> color(n);
    for (int i = 0; i < n; i++){
        cin >> color[i];
    }
    
    // Build tree.
    // We assume that the tree is given as follows:
    // For i = 1 to n-1, input parent of node i (1-indexed).
    // Node 0 is the root.
    vector<vector<int>> adj(n);
    for (int i = 1; i < n; i++){
        int p;
        cin >> p;
        p--; // convert to 0-indexed
        adj[p].push_back(i);
    }
    
    // Euler Tour to flatten the tree.
    // in[u] = first time index of node u in the Euler Tour.
    // out[u] = last time index of node u in the Euler Tour.
    // ET[] will store the color at each position.
    vector<int> in(n), out(n), ET(n);
    int timer = 0;
    function<void(int)> dfs = [&](int v) {
        in[v] = timer;
        ET[timer] = color[v];
        timer++;
        for (int u : adj[v])
            dfs(u);
        out[v] = timer - 1;
    };
    dfs(0);
    
    // Build queries.
    // Each query is given as a node u (1-indexed input),
    // and we want the answer for the subtree of u.
    // That corresponds to the Euler Tour range [in[u], out[u]].
    vector<Query> queries(q);
    for (int i = 0; i < q; i++){
        int u;
        cin >> u;
        u--; // convert to 0-indexed
        queries[i].L = in[u];
        queries[i].R = out[u];
        queries[i].idx = i;
    }
    
    // Mo's algorithm: sort queries by block and then by R.
    int blockSize = static_cast<int>(sqrt(n));
    sort(queries.begin(), queries.end(), [&](const Query &a, const Query &b) {
        if (a.L / blockSize != b.L / blockSize)
            return a.L / blockSize < b.L / blockSize;
        return a.R < b.R;
    });
    
    // Prepare frequency array.
    // Assuming maximum color value is not huge. You can compute max_color if needed.
    int maxColor = *max_element(color.begin(), color.end());
    vector<int> freq(maxColor + 1, 0);
    int distinct = 0; // current number of distinct colors in the window
    
    // Lambdas for add and remove operations.
    auto add = [&](int pos) {
        int c = ET[pos];
        freq[c]++;
        if (freq[c] == 1) distinct++; // new color added
    };
    auto remove = [&](int pos) {
        int c = ET[pos];
        freq[c]--;
        if (freq[c] == 0) distinct--; // color removed completely
    };
    
    // Process queries.
    vector<int> answers(q);
    int curL = 0, curR = -1;
    for (auto &qry : queries) {
        while (curR < qry.R) {
            curR++;
            add(curR);
        }
        while (curR > qry.R) {
            remove(curR);
            curR--;
        }
        while (curL < qry.L) {
            remove(curL);
            curL++;
        }
        while (curL > qry.L) {
            curL--;
            add(curL);
        }
        answers[qry.idx] = distinct;
    }
    
    // Output answers in the original order.
    for (int i = 0; i < q; i++){
        cout << answers[i] << "\n";
    }
    
    return 0;
}

• Subtask 1: 0 / 12
• Subtask 2: 0 / 19
• Subtask 3: 0 / 27
• Subtask 4: 0 / 23
• Subtask 5: 0 / 19