제출 #447222

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
447222		blue	Mergers (JOI19_mergers)	C++17	100 / 100	1267 ms	151476 KiB

mergers

#include <iostream>
#include <set>
#include <vector>
#include <set>
using namespace std;

const int maxN = 500'000;

vector<int> edge[1+maxN];
vector<int> parent(1+maxN, 0);
vector<int> subtree(1+maxN, 1);
vector<int> max_child(1+maxN, -1);

int rt = 1;

void dfs1(int u)
{
    for(int v: edge[u])
    {
        if(parent[u] == v) continue;
        parent[v] = u;
        dfs1(v);
        subtree[u] += subtree[v];
        if(max_child[u] == -1 || subtree[v] > subtree[ max_child[u] ])
            max_child[u] = v;
    }
}




struct state
{
    int i;
    int c;
};

bool operator < (state A, state B)
{
    return A.i < B.i;
}

int S[1+maxN];
vector<int> state_count(1+maxN, 0);

set<state>* T[1+maxN];



void add_state(int u, state x)
{
    // cerr << "add_state " << u << ' ' << x.i << ' ' << x.c << '\n';
    set<state>::iterator it = T[u]->find(x);
    if(it == T[u]->end())
    {
        // cerr << "case A\n";
        if(x.c != state_count[x.i])
            T[u]->insert(x);
    }
    else
    {
        // cerr << "case B\n";
        int ct = it->c + x.c;
        T[u]->erase(it);
        if(ct != state_count[x.i])
        {
            T[u]->insert(state{x.i, ct});
        }
    }
}

vector<int> is_bad(1+maxN, 0);


void dfs2(int u)
{
    // cerr << "dfs2 " << u << ' ' << max_child[u] << '\n';
    if(max_child[u] == -1)
    {
        // cerr << "case 1\n";
        T[u] = new set<state>;
        add_state(u, state{S[u], 1});
    }
    else
    {
        for(int v: edge[u])
        {
            if(v == parent[u]) continue;
            dfs2(v);
        }

        // cerr << "case 2\n";
        T[u] = T[ max_child[u] ];
        // cerr << "x\n";
        add_state(u, state{S[u], 1});
        // cerr << "check X\n";
        for(int v: edge[u])
        {
            // cerr << "v = " << v << '\n';
            if(v == parent[u] || v == max_child[u]) continue;

            for(state s: *(T[v]))
                add_state(u, s);

            T[v]->clear();
        }
    }


    if(T[u]->size() == 0 && u != rt)
    {
        is_bad[u] = 1;
    }
}






vector<int> dsu_parent(1+maxN);
vector<int> dsu_size(1+maxN, 1);

int getRoot(int u)
{
    while(dsu_parent[u] != u)
        u = dsu_parent[u];

    return u;
}

bool dsu_connected(int u, int v)
{
    return getRoot(u) == getRoot(v);
}

void dsu_join(int u, int v)
{
    u = getRoot(u);
    v = getRoot(v);

    if(dsu_size[u] < dsu_size[v]) swap(u, v);

    dsu_parent[v] = u;
    dsu_size[u] += dsu_size[v];
}





int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    subtree[0] = 1e9;

    int N, K;
    cin >> N >> K;

    for(int i = 1; i <= N-1; i++)
    {
        int A, B;
        cin >> A >> B;
        edge[A].push_back(B);
        edge[B].push_back(A);
    }

    for(int i = 1; i <= N; i++)
    {
        cin >> S[i];
        state_count[ S[i] ]++;
    }

    if(N == 2)
    {
        cout << "1\n";
        return 0;
    }

    for(rt = 1; edge[rt].size() == 1; rt++);

    parent[rt] = rt;
    dfs1(rt);
    dfs2(rt);

    for(int i = 1; i <= N; i++) dsu_parent[i] = i;

    for(int i = 1; i <= N; i++)
    {
        if(i == rt) continue;
        if(!is_bad[i]) dsu_join(i, parent[i]);
    }

    set<int> edg[1+N];
    for(int i = 1; i <= N; i++)
    {
        for(int j: edge[i])
        {
            if(!dsu_connected(i, j))
            {
                edg[ getRoot(i) ].insert( getRoot(j) );
            }
        }
    }

    int cnt = 0;
    for(int i = 1; i <= N; i++)
    {
        if(getRoot(i) == i && edg[i].size() == 1)
            cnt++;
    }

    cout << (cnt+1)/2 << '\n';
}

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