Submission #1143688

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

joi2019_ho_t5

#include <bits/stdc++.h>
using namespace std;
 
// -------------------- Data Structure -------------------------
// For a given distance d (from some “root”) we store:
//   cnt  = number of nodes at that distance,
//   cand = if cnt==1 then the candidate speciality (otherwise -1).
struct PairEntry {
    int cnt;
    int cand;
};
 
// Combine two PairEntry values for the same distance.
inline PairEntry combinePair(const PairEntry &a, const PairEntry &b) {
    PairEntry res;
    res.cnt = a.cnt + b.cnt;
    if(res.cnt == 1)
        res.cand = (a.cnt == 1 ? a.cand : b.cand);
    else
        res.cand = -1;
    return res;
}
 
// Combine two full distributions A and B (element–wise).
// If one vector is shorter, missing slots are treated as {0, -1}.
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++){
        PairEntry a = (i < A.size() ? A[i] : PairEntry{0, -1});
        PairEntry b = (i < B.size() ? B[i] : PairEntry{0, -1});
        res[i] = combinePair(a, b);
    }
    return res;
}
 
// When merging a child’s distribution into a parent’s distribution,
// we must “shift” the child’s distribution by 1.
void mergeShift(vector<PairEntry>& base, const vector<PairEntry>& child) {
    int shift = 1;
    int newSize = max((int)base.size(), (int)child.size() + shift);
    if(base.size() < newSize)
        base.resize(newSize, {0, -1});
    for (int i = 0; i < child.size(); i++){
        int pos = i + shift;
        base[pos] = combinePair(base[pos], child[i]);
    }
}
 
// --------------------- Global Variables ---------------------
int N, M;
vector<int> C;               // Speciality for each city (0-indexed)
vector<vector<int>> adj;     // Tree: adj[u] holds neighbors of u.
 
// dp[x] will store the distribution for the subtree rooted at x.
// At index d, dp[x][d] = { count, candidate } for nodes in x's subtree at distance d (from x).
// (Note: In the worst–case a chain would yield size ~O(N). Using small–to–large merging
//  ensures that the total number of element copies is only O(N log N) rather than O(N^2).)
vector<vector<PairEntry>> dp;
 
// ans[x] will be the final answer for city x.
vector<int> ans;
 
// --------------------- DFS 1: Down–DFS with Small–to–Large Merging -------------------------
// Compute dp[x] for the subtree rooted at x.
// Initialize dp[x] with {1, C[x]} at distance 0. Then for each child,
// merge its shifted dp (i.e. every distance increased by 1) into dp[x].
// The small–to–large trick is used so that when merging dp[v] (shifted) into dp[x],
// we swap them if the child’s dp (after shifting) is larger than dp[x].
void dfs_down(int x, int parent) {
    dp[x].clear();
    dp[x].push_back({1, C[x]}); // distance 0: city x.
    
    for (int v : adj[x]) {
        if(v == parent) continue;
        dfs_down(v, x);
        // Create a temporary vector for dp[v] shifted by 1.
        vector<PairEntry> temp(dp[v].size() + 1, {0, -1});
        for (int i = 0; i < dp[v].size(); i++){
            temp[i+1] = dp[v][i];
        }
        // Small–to–large: if temp is larger than dp[x], swap them.
        if(temp.size() > dp[x].size())
            swap(dp[x], temp);
        // Merge temp into dp[x] element–wise.
        for (int i = 0; i < temp.size(); i++){
            if(i < dp[x].size()){
                dp[x][i] = combinePair(dp[x][i], temp[i]);
            } else {
                dp[x].push_back(temp[i]);
            }
        }
    }
}
 
// --------------------- DFS 2: Up–DFS (Re–rooting) -------------------------
// For node x, given an “up” distribution (contributions from nodes outside x’s subtree),
// compute the full distribution F[x] = combineDist(up, dp[x]) and then count the number
// of distinct specialities among layers (d>=1) that are unique.
void dfs_up(int x, int parent, const vector<PairEntry>& up) {
    // Full distribution for x.
    vector<PairEntry> full = combineDist(up, dp[x]);
    
    // For every distance d>=1, if full[d].cnt==1 then record its candidate 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 compute an up–distribution for each child.
    vector<int> children;
    for (int v : adj[x])
        if(v != parent)
            children.push_back(v);
    int k = children.size();
    
    // Build prefix and suffix arrays for the contributions from x’s children.
    // Each child’s branch contribution is dp[v] shifted by 1.
    vector< vector<PairEntry> > pref(k+1), suf(k+1);
    pref[0].clear();
    for (int i = 0; i < k; i++){
        int v = children[i];
        vector<PairEntry> branch;
        branch.resize(dp[v].size() + 1, {0, -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;
        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]);
    }
    
    // For each child v, compute its up–distribution and continue DFS.
    for (int i = 0; i < k; i++){
        int v = children[i];
        // x’s own contribution: one node (x) at distance 0.
        vector<PairEntry> baseVec;
        baseVec.push_back({1, C[x]});
        // Siblings’ contribution: merge contributions from children before and after v.
        vector<PairEntry> siblings = combineDist(pref[i], suf[i+1]);
        // R = (x itself) + (siblings).
        vector<PairEntry> R = combineDist(baseVec, siblings);
        // Combine with parent's up–distribution.
        vector<PairEntry> newUp = combineDist(up, R);
        // Shift newUp by 1 for child v.
        vector<PairEntry> shiftedNewUp;
        shiftedNewUp.resize(newUp.size()+1, {0,-1});
        for (int j = 0; j < (int)newUp.size(); j++){
            shiftedNewUp[j+1] = newUp[j];
        }
        dfs_up(v, x, shiftedNewUp);
    }
}
 
// ------------------------- Main -------------------------
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--; // Convert to 0-index.
        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 the dp distribution for each subtree (using small-to-large merging
    // to keep memory usage under control).
    dfs_down(0, -1);
    
    // Second DFS: re–rooting.
    vector<PairEntry> up; // For the root there is no parent's contribution.
    up.clear();
    dfs_up(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