Submission #1050853

#	Time	Username	Problem	Language	Result	Execution time	Memory
1050853		Jarif_Rahman	Mergers (JOI19_mergers)	C++17	34 / 100	3052 ms	14960 KiB

mergers

#include <bits/stdc++.h>
#define pb push_back
#define f first
#define sc second
using namespace std;
typedef long long int ll;
typedef string str;

int n, k;
vector<vector<int>> tree;
vector<int> sz;
vector<int> state, scnt;
vector<bool> sep;
int sep_cnt = 0;

void add(int nd, int ss, vector<int> &cnt){
    for(int x: tree[nd]) if(x != ss) add(x, nd, cnt);
    cnt[state[nd]]++;
}
void dfs(int nd, int ss){
    for(int x: tree[nd]) if(x != ss) dfs(x, nd);
    vector<int> cnt(k, 0);
    add(nd, ss, cnt);
    for(int i = 0; i < k; i++) sep[nd] = sep[nd]&(cnt[i] == 0 || cnt[i] == scnt[i]);
    if(sep[nd]) sep_cnt++;
}

int ans = 0;
bool exta = 0;
int dfs2(int nd, int ss, int sep_cnt_ = 0){
    int c = 0;
    for(int x: tree[nd]) if(x != ss) c+=dfs2(x, nd, sep_cnt_+sep[nd]);
    if(ss == -1) return 0;
    if(!sep[nd]) return c;
    c++;
    if(c == 1) ans++;
    else if(c+sep_cnt_ == sep_cnt) exta = 1;
    return c;
}

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

    cin >> n >> k;
    tree.resize(n);
    state.resize(n);
    scnt.resize(k, 0);
    sep.resize(n, 1);

    for(int i = 0; i < n-1; i++){
        int a, b; cin >> a >> b; a--, b--;
        tree[a].push_back(b);
        tree[b].push_back(a);
    }
    for(int i = 0; i < n; i++){
        cin >> state[i]; state[i]--;
        scnt[state[i]]++;
    }

    dfs(0, -1);
    dfs2(0, -1);
    ans+=exta;
    ans++;
    ans/=2;
    cout << ans << "\n";
}

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