제출 #333825

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
333825		12tqian	Mergers (JOI19_mergers)	C++17	100 / 100	1995 ms	193352 KiB

mergers

#include <bits/stdc++.h>
struct LCAJump {
    int n;
    std::vector<std::vector<int>> par;
    std::vector<std::vector<int>> adj;
    std::vector<int> depth;
    void init(int _n) {
        n = _n;
        int d = 1;
        while ((1 << d) < n) d++;
        par.assign(d, std::vector<int>(n));
        adj.resize(n);
        depth.resize(n);
    }
    void ae(int x, int y) {
        adj[x].push_back(y);
        adj[y].push_back(x);
    }
    void gen(int root = 0) {
        par[0][root] = root;
        dfs(root);
    }
    void dfs(int src = 0) {
        for (int i = 1; i < par.size(); i++) {
            par[i][src] = par[i - 1][par[i - 1][src]];
        }
        for (int nxt: adj[src]) {
            if (nxt == par[0][src]) continue;
            depth[nxt] = depth[par[0][nxt] = src] + 1;
            dfs(nxt);
        }
    }
    int jump(int x, int d) {
        for (int i = 0; i < par.size(); i++) {
            if ((d >> i) & 1) {
                x = par[i][x];
            }
        }
        return x;
    }
    int lca(int x, int y) {
        if (depth[x] < depth[y]) std::swap(x, y);
        x = jump(x, depth[x] - depth[y]);
        if (x == y) return x;
        for (int i = par.size() - 1; i >= 0; i--) {
            int nx = par[i][x];
            int ny = par[i][y];
            if (nx != ny) x = nx, y = ny;
        }
        return par[0][x];
    }
};

int main() {
    using namespace std;
    ios_base::sync_with_stdio(0);
    int n, k;
    cin >> n >> k;
    LCAJump L;
    L.init(n);
    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);
        L.ae(u, v);
    }
    L.gen(0);
    vector<vector<int>> loc(k);
    vector<int> s(n);
    for (int i = 0; i < n; i++) 
        cin >> s[i], s[i]--, loc[s[i]].push_back(i);
    if (n == 1) {
        cout << 0 << '\n';
        return 0;
    }
    vector<int> add(n, 1);
    for (int i = 0; i < k; i++) {
        if (loc[i].empty())
            continue;
        int lca = loc[i].back();
        for (int x : loc[i])
            lca = L.lca(lca, x);
        add[lca] -= (int) loc[i].size();
    }
    vector<int> tot(n);
    for (int x : s)
        tot[x]++;
    vector<multiset<int>> store(n);
    vector<int> id(n);
    iota(id.begin(), id.end(), 0);
    vector<int> bad(n);
    function<int(int, int)> combine = [&](int src, int par) -> int {
        int res = add[src];
        for (int nxt : adj[src]){ 
            if (nxt == par)
                continue;
            res += combine(nxt, src);
        }
        if (res == 0 && src != 0)
            bad[src] = 1;
        return res;
    };
    combine(0, -1);
    int amt = 0;
    for (int i = 1; i < n; i++)
        amt += bad[i];
    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;
}

컴파일 시 표준 에러 (stderr) 메시지

mergers.cpp: In member function 'void LCAJump::dfs(int)':
mergers.cpp:24:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   24 |         for (int i = 1; i < par.size(); i++) {
      |                         ~~^~~~~~~~~~~~
mergers.cpp: In member function 'int LCAJump::jump(int, int)':
mergers.cpp:34:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   34 |         for (int i = 0; i < par.size(); i++) {
      |                         ~~^~~~~~~~~~~~

• Subtask 1: 10 / 10
• Subtask 2: 24 / 24
• Subtask 3: 14 / 14
• Subtask 4: 22 / 22
• Subtask 5: 30 / 30