Submission #1143684

#	Time	Username	Problem	Language	Result	Execution time	Memory
1143684		Otalp	Unique Cities (JOI19_ho_t5)	C++20	4 / 100	550 ms	339968 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
using namespace std;

// Structure to hold (count, candidate speciality) for a given distance.
struct PairEntry {
    int cnt;
    int cand; // if cnt==1, then cand is the speciality of the unique city; otherwise, cand is -1.
};

// Combine two distributions A and B (element–wise) so that for each index i the new
// count is A[i].cnt + B[i].cnt and if that sum is 1 the candidate is the one that appears.
vector<PairEntry> combineDist(const vector<PairEntry>& A, const vector<PairEntry>& B) {
    int n = max(A.size(), B.size());
    vector<PairEntry> res(n, {0, -1});
    for (int i = 0; i < n; i++){
        int a_cnt = (i < (int)A.size() ? A[i].cnt : 0);
        int b_cnt = (i < (int)B.size() ? B[i].cnt : 0);
        int tot = a_cnt + b_cnt;
        int cand = -1;
        if(tot == 1) {
            if(i < (int)A.size() && A[i].cnt == 1)
                cand = A[i].cand;
            else if(i < (int)B.size() && B[i].cnt == 1)
                cand = B[i].cand;
        }
        res[i] = {tot, cand};
    }
    return res;
}
 
// Merge a distribution 'child' into 'base', but with a shift of 1 (because if we attach child’s subtree
// then every node in that subtree is one road farther from the parent).
void mergeShift(vector<PairEntry>& base, const vector<PairEntry>& child) {
    int shift = 1;
    int newSize = max((int)base.size(), (int)child.size() + shift);
    base.resize(newSize, {0, -1});
    for (int i = 0; i < (int)child.size(); i++){
        int pos = i + shift;
        int tot = base[pos].cnt + child[i].cnt;
        int cand = -1;
        if(tot == 1) {
            if(base[pos].cnt == 1)
                cand = base[pos].cand;
            else
                cand = child[i].cand;
        }
        base[pos].cnt = tot;
        base[pos].cand = cand;
    }
}
 
// Global variables.
int N, M;
vector<int> C;               // Speciality for each city (0-indexed).
vector<vector<int>> adj;     // Tree (adjacency list).
vector<vector<PairEntry>> dp; // dp[x] will store the distribution for the subtree rooted at x.
vector<int> ans;             // ans[x] will be the final answer for city x.
 
// dfs1 computes dp[x] = distribution for the subtree rooted at x.
// At distance 0 we record city x itself.
void dfs1(int x, int parent) {
    dp[x].clear();
    dp[x].push_back({1, C[x]}); // distance 0: only x itself.
    for (int v : adj[x]) {
        if(v == parent) continue;
        dfs1(v, x);
        mergeShift(dp[x], dp[v]);
    }
}
 
// dfs2 does the re–rooting. For node x we are given 'up' – the distribution of all nodes that are not in x’s subtree,
// as seen from x. Then the full distribution for x is F[x] = combineDist(up, dp[x]).
// From F[x] we can count the number of distances d>=1 that contain a unique city and then compute the number of distinct specialities.
void dfs2(int x, int parent, const vector<PairEntry>& up) {
    // Full distribution F[x] = (up from parent's side) + (dp[x] from x’s subtree)
    vector<PairEntry> full = combineDist(up, dp[x]);
    
    // For every distance d>=1, if full[d].cnt==1 then the unique city at that distance contributes its speciality.
    unordered_set<int> uniqSpecs;
    for (int d = 1; d < (int)full.size(); d++){
        if(full[d].cnt == 1)
            uniqSpecs.insert(full[d].cand);
    }
    ans[x] = uniqSpecs.size();
    
    // For re–rooting we want to compute for each child v of x a new up–distribution.
    // Let children[] be the list of x’s children (neighbors except the parent).
    vector<int> children;
    for (int v : adj[x])
        if(v != parent)
            children.push_back(v);
    int k = children.size();
    
    // Build prefix and suffix arrays over the children.
    // For each child, the “branch” distribution is dp[v] shifted by 1.
    vector< vector<PairEntry> > pref(k+1), suf(k+1);
    pref[0].clear(); // identity: an empty distribution (which we interpret as all zeros).
    for (int i = 0; i < k; i++){
        int v = children[i];
        vector<PairEntry> branch;
        if(dp[v].empty()){
            branch.clear();
        } else {
            branch.resize(dp[v].size() + 1, {0, -1});
            // Shift dp[v] by 1.
            for (int j = 0; j < (int)dp[v].size(); j++){
                branch[j+1] = dp[v][j];
            }
        }
        pref[i+1] = combineDist(pref[i], branch);
    }
    suf[k].clear();
    for (int i = k-1; i >= 0; i--){
        int v = children[i];
        vector<PairEntry> branch;
        if(dp[v].empty()){
            branch.clear();
        } else {
            branch.resize(dp[v].size() + 1, {0, -1});
            for (int j = 0; j < (int)dp[v].size(); j++){
                branch[j+1] = dp[v][j];
            }
        }
        suf[i] = combineDist(branch, suf[i+1]);
    }
    
    // Now, for each child v we want to compute its up–distribution.
    // The idea is that the parent's full distribution F[x] comes from:
    //   (i) x itself (at distance 0), plus
    //   (ii) contributions from all children.
    // For child v, we remove its branch from F[x] and then “shift” the result by 1.
    for (int i = 0; i < k; i++){
        int v = children[i];
        // x’s own contribution is a base distribution: one node (x) at distance 0.
        vector<PairEntry> baseVec;
        baseVec.push_back({1, C[x]});
        // Combine contributions from siblings: merge (prefix from children[0..i-1] and suffix from children[i+1..k-1]).
        vector<PairEntry> siblings = combineDist(pref[i], suf[i+1]);
        // Combine x’s own contribution with siblings.
        vector<PairEntry> R = combineDist(baseVec, siblings);
        // Also, include the parent's up–distribution that was passed to x.
        vector<PairEntry> newUp = combineDist(up, R);
        // Now, shift newUp by 1 to get the up–distribution as seen from child v.
        vector<PairEntry> shiftedNewUp;
        if(newUp.empty()){
            shiftedNewUp.clear();
        } else {
            shiftedNewUp.resize(newUp.size() + 1, {0, -1});
            for (int j = 0; j < (int)newUp.size(); j++){
                shiftedNewUp[j+1] = newUp[j];
            }
        }
        dfs2(v, x, shiftedNewUp);
    }
}
 
// Main reads the tree and specialities, computes dp and answers, then prints them.
int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    
    cin >> N >> M;
    adj.assign(N, vector<int>());
    dp.assign(N, vector<PairEntry>());
    ans.assign(N, 0);
    C.resize(N);
    
    // Read the N-1 roads.
    for (int i = 0; i < N-1; i++){
        int u, v;
        cin >> u >> v;
        u--; v--; // converting to 0-indexed.
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    // Read the speciality for each city.
    for (int i = 0; i < N; i++){
        cin >> C[i];
    }
    
    // First DFS: compute dp for each subtree.
    dfs1(0, -1);
    
    // For the root, there is no parent's contribution (up is empty).
    vector<PairEntry> up;
    up.clear();
    dfs2(0, -1, up);
    
    // Output the answer for each city.
    for (int i = 0; i < N; i++){
        cout << ans[i] << "\n";
    }
    
    return 0;
}

• Subtask 1: 4 / 4
• Subtask 2: 0 / 32
• Subtask 3: 0 / 32
• Subtask 4: 0 / 32