Submission #1312969

#	Time	Username	Problem	Language	Result	Execution time	Memory
1312969		4QT0R	Sumtree (INOI20_sumtree)	C++20	100 / 100	502 ms	65972 KiB

Main

// gemini code for testing
#include <iostream>
#include <vector>
#include <algorithm>
#include <random>

using namespace std;

// --- Constants ---
const int MOD = 1e9 + 7;
const int MAXN = 200005;
// Max value for combinations is R + N approx 3e5 + 2e5 = 5e5
const int MAX_FACT = 600005; 

// --- Math Helpers ---
long long fact[MAX_FACT], invFact[MAX_FACT];

long long power(long long base, long long exp) {
    long long res = 1;
    base %= MOD;
    while (exp > 0) {
        if (exp % 2 == 1) res = (res * base) % MOD;
        base = (base * base) % MOD;
        exp /= 2;
    }
    return res;
}

long long modInverse(long long n) {
    return power(n, MOD - 2);
}

void precomputeFactorials() {
    fact[0] = 1;
    invFact[0] = 1;
    for (int i = 1; i < MAX_FACT; i++) {
        fact[i] = (fact[i - 1] * i) % MOD;
    }
    invFact[MAX_FACT - 1] = modInverse(fact[MAX_FACT - 1]);
    for (int i = MAX_FACT - 2; i >= 0; i--) {
        invFact[i] = (invFact[i + 1] * (i + 1)) % MOD;
    }
}

long long nCr(int n, int r) {
    if (r < 0 || r > n) return 0;
    return (((fact[n] * invFact[r]) % MOD) * invFact[n - r]) % MOD;
}

// --- Tree Data ---
vector<int> adj[MAXN];
int tin[MAXN], tout[MAXN], sub_sz[MAXN], depth[MAXN], parent[MAXN];
int timer;
int n, r_initial;

void dfs(int u, int p, int d) {
    tin[u] = ++timer;
    sub_sz[u] = 1;
    depth[u] = d;
    parent[u] = p;
    for (int v : adj[u]) {
        if (v != p) {
            dfs(v, u, d + 1);
            sub_sz[u] += sub_sz[v];
        }
    }
    tout[u] = timer;
}

// --- Treap (for managing children lists) ---
struct TreapNode {
    int id; // Vertex ID
    int priority;
    int key; // Uses tin[id] as key
    int l, r;
    
    // Aggregates
    long long sum_val; // Sum of tag values
    long long sum_sz;  // Sum of subtree sizes
    
    // Values of this specific node
    long long self_val; 
    long long self_sz; 
} tree[MAXN];

int root_treap[MAXN]; // The Treap Root for each fixed node
mt19937 rng(1337);

void update_treap_node(int t) {
    if (!t) return;
    tree[t].sum_val = tree[t].self_val;
    tree[t].sum_sz = tree[t].self_sz;
    if (tree[t].l) {
        tree[t].sum_val += tree[tree[t].l].sum_val;
        tree[t].sum_sz += tree[tree[t].l].sum_sz;
    }
    if (tree[t].r) {
        tree[t].sum_val += tree[tree[t].r].sum_val;
        tree[t].sum_sz += tree[tree[t].r].sum_sz;
    }
}

void split(int t, int key, int &l, int &r) {
    if (!t) {
        l = r = 0;
        return;
    }
    if (tree[t].key <= key) {
        split(tree[t].r, key, tree[t].r, r);
        l = t;
    } else {
        split(tree[t].l, key, l, tree[t].l);
        r = t;
    }
    update_treap_node(t);
}

void merge(int &t, int l, int r) {
    if (!l || !r) t = l ? l : r;
    else if (tree[l].priority > tree[r].priority) {
        merge(tree[l].r, tree[l].r, r);
        t = l;
    } else {
        merge(tree[r].l, l, tree[r].l);
        t = r;
    }
    update_treap_node(t);
}

// --- HLD + SegTree (for Nearest Fixed Ancestor) ---
int head[MAXN], pos[MAXN], hld_sz[MAXN];
pair<int, int> seg[4 * MAXN]; // Stores {depth, node_index}
int cur_pos = 0;

void dfs_hld(int u) {
    hld_sz[u] = 1;
    int max_sz = 0, big_child = -1;
    for (int &v : adj[u]) {
        if (v != parent[u]) {
            dfs_hld(v);
            hld_sz[u] += hld_sz[v];
            if (hld_sz[v] > max_sz) {
                max_sz = hld_sz[v];
                big_child = v;
            }
        }
    }
    // Optimization: move heavy child to front
    // (Omitted for brevity, standard HLD optimization)
}

void decompose(int u, int h) {
    head[u] = h;
    pos[u] = ++cur_pos;
    
    int big_child = -1, max_sz = -1;
    for (int v : adj[u]) {
        if (v != parent[u] && hld_sz[v] > max_sz) {
            max_sz = hld_sz[v];
            big_child = v;
        }
    }
    if (big_child != -1) decompose(big_child, h);
    for (int v : adj[u]) {
        if (v != parent[u] && v != big_child) {
            decompose(v, v);
        }
    }
}

void update_seg(int node, int start, int end, int idx, pair<int, int> val) {
    if (start == end) {
        seg[node] = val;
        return;
    }
    int mid = (start + end) / 2;
    if (idx <= mid) update_seg(node * 2, start, mid, idx, val);
    else update_seg(node * 2 + 1, mid + 1, end, idx, val);
    seg[node] = max(seg[node * 2], seg[node * 2 + 1]);
}

pair<int, int> query_seg(int node, int start, int end, int l, int r) {
    if (r < start || end < l) return {-1, -1};
    if (l <= start && end <= r) return seg[node];
    int mid = (start + end) / 2;
    return max(query_seg(node * 2, start, mid, l, r), query_seg(node * 2 + 1, mid + 1, end, l, r));
}

int get_nearest_fixed_ancestor(int u) {
    pair<int, int> res = {-1, -1};
    while (u != -1) { // -1 when we go above root, but loop condition handles head[u] logic
        res = max(res, query_seg(1, 1, n, pos[head[u]], pos[u]));
        if (head[u] == 1) break; // Reached root chain
        u = parent[head[u]];
    }
    return res.second;
}

// --- Global Answer Management ---
long long current_ans = 1;
int zero_count = 0;
long long node_ways[MAXN]; 
bool is_fixed[MAXN];
long long fixed_val[MAXN];

long long calculate_ways(int u) {
    long long load = 0;
    long long size_deduction = 0;
    if (root_treap[u]) {
        load = tree[root_treap[u]].sum_val;
        size_deduction = tree[root_treap[u]].sum_sz;
    }
    
    long long rem = fixed_val[u] - load;
    long long scope_size = sub_sz[u] - size_deduction;
    
    if (rem < 0) return 0;
    // Stars and bars: (Slack + ScopeSize - 1) choose (ScopeSize - 1)
    return nCr(rem + scope_size - 1, scope_size - 1);
}

void update_global_ans(int u, bool adding) {
    long long w = node_ways[u];
    if (!adding) {
        if (w == 0) zero_count--;
        else current_ans = (current_ans * modInverse(w)) % MOD;
    } else {
        w = calculate_ways(u);
        node_ways[u] = w;
        if (w == 0) zero_count++;
        else current_ans = (current_ans * w) % MOD;
    }
}

// --- Main ---
int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    precomputeFactorials();
    
    if (!(cin >> n >> r_initial)) return 0;
    
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    
    dfs(1, -1, 0);
    dfs_hld(1);
    decompose(1, 1);
    
    // Init Segment Tree
    for(int i=0; i<4*MAXN; i++) seg[i] = {-1, -1};
    
    // Init Treap Nodes
    for(int i=1; i<=n; i++) {
        tree[i].id = i;
        tree[i].priority = rng();
        tree[i].key = tin[i];
        tree[i].l = tree[i].r = 0;
        tree[i].self_val = 0;
        tree[i].self_sz = sub_sz[i]; // Initial full subtree size
        update_treap_node(i);
    }
    
    // Set up Root
    is_fixed[1] = true;
    fixed_val[1] = r_initial;
    tree[1].self_val = r_initial;
    update_treap_node(1);
    
    update_seg(1, 1, n, pos[1], {depth[1], 1});
    update_global_ans(1, true); 
    
    cout << (zero_count > 0 ? 0 : current_ans) << "\n";
    
    int q;
    cin >> q;
    while(q--) {
        int type;
        cin >> type;
        if (type == 1) { // Add Tag
            int u, val;
            cin >> u >> val;
            
            // 1. Find Ancestor
            int anc = get_nearest_fixed_ancestor(u);
            
            // 2. Remove Anc from answer
            update_global_ans(anc, false);
            
            // 3. Extract children from Anc that belong to U
            // U's subtree range is [tin[u], tout[u]]
            // Treap is sorted by tin. Just split by range.
            int t_anc = root_treap[anc];
            int l, m, r_part;
            
            split(t_anc, tout[u], m, r_part); // m has keys <= tout[u]
            split(m, tin[u]-1, l, m);         // m has keys >= tin[u]
            // Now 'm' contains exactly the descendants of u
            
            root_treap[u] = m;
            
            // 4. Setup U
            is_fixed[u] = true;
            fixed_val[u] = val;
            tree[u].self_val = val;
            // The treap node for U stores its RAW subtree size. 
            // The deduction happens when calculating 'calculate_ways', not in the node struct itself.
            tree[u].self_sz = sub_sz[u]; 
            update_treap_node(u);
            
            // 5. Insert U into Anc's treap
            // Reassemble anc: l + u + r_part
            int temp;
            merge(temp, l, u); // u is index in tree[] array
            merge(root_treap[anc], temp, r_part);
            
            // 6. Update SegTree
            update_seg(1, 1, n, pos[u], {depth[u], u});
            
            // 7. Update Global Answers
            update_global_ans(u, true);
            update_global_ans(anc, true);
            
        } else { // Remove Tag
            int u;
            cin >> u;
            
            // 1. Get Parent in Fixed Tree
            // Temporarily unmark u to find its ancestor
            update_seg(1, 1, n, pos[u], {-1, -1});
            int anc = get_nearest_fixed_ancestor(u);
            
            // 2. Remove U and Anc from answer
            // Note: we must remove U from global calc BEFORE unsetting its data
            update_global_ans(u, false); 
            update_global_ans(anc, false);
            
            // 3. Remove U from Anc's treap
            // U is a direct child in Anc's treap. Isolate it.
            int t_anc = root_treap[anc];
            int l, mid, r_part;
            
            split(t_anc, tin[u], mid, r_part); // mid <= tin[u]
            split(mid, tin[u]-1, l, mid);      // mid == u (since key is unique)
            
            // 4. Merge U's children back into Anc
            // U's children are in root_treap[u].
            // We merge: l + root_treap[u] + r_part
            int temp;
            merge(temp, l, root_treap[u]);
            merge(root_treap[anc], temp, r_part);
            
            // 5. Finalize State
            is_fixed[u] = false;
            // SegTree already updated in step 1
            
            // 6. Recalculate Anc
            update_global_ans(anc, true);
        }
        
        cout << (zero_count > 0 ? 0 : current_ans) << "\n";
    }
    
    return 0;
}

• Subtask 1: 10 / 10
• Subtask 2: 20 / 20
• Subtask 3: 20 / 20
• Subtask 4: 20 / 20
• Subtask 5: 30 / 30