Submission #333575

#	Time	Username	Problem	Language	Result	Execution time	Memory
333575		12tqian	Long Mansion (JOI17_long_mansion)	C++17	0 / 100	248 ms	262148 KiB

long_mansion

#include <bits/stdc++.h>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>

using namespace std;
using namespace __gnu_pbds;
template <class T> using Tree = tree<T, null_type, less<T>, rb_tree_tag,tree_order_statistics_node_update>;

template <class T> struct LazySeg {
    std::vector<T> sum, lazy;
    int sz;
    void init(int sz_) {
        sz = 1;
        while (sz < sz_) sz *= 2;
        sum.assign(2 * sz, 0);
        lazy.assign(2 * sz, 0);
    }
    void push(int ind, int L, int R) {
        sum[ind] += (R - L + 1) * lazy[ind];
        if (L != R) {
            lazy[2 * ind] += lazy[ind];
            lazy[2 * ind + 1] += lazy[ind];
        }
        lazy[ind] = 0;
    }
    void pull(int ind) {
        sum[ind] = sum[2 * ind] + sum[2 * ind + 1];
    }
    void build() {
        for (int i = sz - 1; i >= 1; i--) {
            pull(i);
        }
    }
    void upd(int lo, int hi, T inc, int ind = 1, int L = 0, int R = -1) {
        if (R == -1) R += sz;
        push(ind, L, R);
        if (hi < L || R < lo) return ;
        if (lo <= L && R <= hi) {
            lazy[ind] = inc;
            push(ind, L, R);
            return;
        }
        int M = (L + R) / 2;
        upd(lo, hi, inc, 2 * ind, L, M);
        upd(lo, hi, inc, 2 * ind + 1, M + 1, R);
        pull(ind);
    }
    T qsum(int lo, int hi, int ind = 1, int L = 0, int R = -1) {
        if (R == -1) R += sz;
        push(ind, L, R);
        if (lo > R || L > hi) return 0;
        if (lo <= L && R <= hi) return sum[ind];
        int M = (L + R) / 2;
        return qsum(lo, hi, 2 * ind, L, M) + qsum(lo, hi, 2 * ind + 1, M + 1, R);
    }
};

int main() {
    // ios_base::sync_with_stdio(0);
    int n, k;
    cin >> n >> k;
    vector<vector<int>> adj(n);
    for (int i = 0; i < n - 1; i++) {
        int u, v; cin >> u >> v;
        u--, v--;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    vector<int> s(n);
    for (int i = 0; i < n; i++) 
        cin >> s[i], s[i]--;
    if (n == 1) {
        cout << 1 << '\n';
        return 0;
    }
    vector<int> L(n), R(n), parent(n);
    int ti = 0;
    function<int(int, int)> dfs_precomp = [&](int src, int par) -> int {
        int mx = ti++;
        parent[src] = par;
        L[src] = mx;
        for (int nxt : adj[src]) {
            if (nxt == par) 
                continue;
            mx = max(mx, dfs_precomp(nxt, src));
        }
        R[src] = mx;
        return mx;
    };
    dfs_precomp(0, -1);
    LazySeg<int> seg;
    seg.init(n);
    vector<int> oc(n), tot(n);
    vector<Tree<int>> num(n);
    for (int i = 0; i < n; i++)
        num[s[i]].insert(L[i]);
    for (int x : s)
        tot[x]++;
    function<void(int, int)> dfs_edges = [&](int src, int par) {
        for (int nxt : adj[src]) {
            if (nxt == par)
                continue;
            dfs_edges(nxt, src);
        }
        int val = num[s[src]].order_of_key(R[src] + 1);
        val -= num[s[src]].order_of_key(L[src]);
        if (val == 1) {
            // update the subtree    
            seg.upd(L[src], R[src], 1);
            seg.upd(L[src], L[src], -1);
        } 
        if (val == tot[s[src]]) {
            // update above
            seg.upd(L[0], R[0], 1);
            seg.upd(L[src], R[src], -1);
            seg.upd(L[src], L[src], 1);
        }
    };
    dfs_edges(0, -1);
    vector<int> bad(n);
    int amt = 0;
    for (int i = 1; i < n; i++) 
        if (seg.qsum(L[i], L[i]) == k) 
            bad[i] = 1, amt++;
    int leaves = 0;
    function<int(int, int)> dfs_sub = [&](int src, int par) -> int {
        int sub = bad[src];
        for (int nxt : adj[src]) {
            if (nxt == par)
                continue;
            sub += dfs_sub(nxt, src);
        }
        if (sub == 1 && bad[src]) 
            leaves++;
        else if (sub == amt && bad[src])
            leaves++;
        return sub;
    };
    dfs_sub(0, -1);
    int ans = (leaves + 1) / 2;
    cout << ans << '\n';
    return 0;
}

• Subtask 1: 0 / 5
• Subtask 2: 0 / 5
• Subtask 3: 0 / 15
• Subtask 4: 0 / 75