Submission #227202

#	Time	Username	Problem	Language	Result	Execution time	Memory
227202		MKopchev	Mergers (JOI19_mergers)	C++14	100 / 100	1176 ms	75176 KiB

mergers

#include <bits/stdc++.h>
using namespace std;

const int nmax=5e5+42;

int n,k;

int colour[nmax],cnt[nmax];

vector<int> adj[nmax];

int wrong=0;

int seen[nmax];

void add(int val)
{
    val=colour[val];

    if(seen[val]==0)wrong++;
    seen[val]++;
    if(seen[val]==cnt[val])wrong--;
}

void sub(int val)
{
    val=colour[val];

    if(seen[val]==cnt[val])wrong++;
    seen[val]--;
    if(seen[val]==0)wrong--;
}

int SZ[nmax];

void dfs_sz(int node,int par)
{
    SZ[node]=1;
    for(auto k:adj[node])
        if(k!=par)
        {
            dfs_sz(k,node);
            SZ[node]+=SZ[k];
        }
}

int can_divide=0;

void dfs_clean(int node,int par)
{
    sub(node);

    for(auto k:adj[node])
        if(k!=par)
        {
            dfs_clean(k,node);
        }
}

void dfs_add(int node,int par)
{
    add(node);

    for(auto k:adj[node])
        if(k!=par)
        {
            dfs_add(k,node);
        }
}

int ROOT;
int dfs(int node,int par)
{
    int found=0;

    int big=0;

    for(auto k:adj[node])
        if(k!=par&&SZ[big]<SZ[k])big=k;

    for(auto k:adj[node])
        if(k!=par&&big!=k)
        {
            found=dfs(k,node)+found;

            dfs_clean(k,node);
        }
    /*
    cout<<node<<" -> "<<big<<" "<<found<<" "<<wrong<<endl;

    cout<<"seen(should be empty): ";for(int i=1;i<=k;i++)cout<<seen[i]<<" ";cout<<endl;
    */
    if(big)found=found+dfs(big,node);


    for(auto k:adj[node])
        if(k!=par&&big!=k)
            dfs_add(k,node);

    add(node);

    //cout<<node<<" -> "<<big<<" "<<found<<" "<<wrong<<endl;

    //cout<<"seen: ";for(int i=1;i<=k;i++)cout<<seen[i]<<" ";cout<<endl;

    if(wrong==0)
    {
        if(node!=ROOT)found++;

        if(found==1)can_divide++;

        return 1;
    }

    return found;
}
int main()
{
    scanf("%i%i",&n,&k);

    int u,v;
    for(int i=1;i<n;i++)
    {
        scanf("%i%i",&u,&v);
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    for(int i=1;i<=n;i++)
    {
        scanf("%i",&colour[i]);
        cnt[colour[i]]++;
    }

    ROOT=1;

    dfs_sz(ROOT,0);

    dfs(ROOT,0);

    printf("%i\n",(can_divide+1)/2);

    return 0;
}

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:119:10: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
     scanf("%i%i",&n,&k);
     ~~~~~^~~~~~~~~~~~~~
mergers.cpp:124:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
         scanf("%i%i",&u,&v);
         ~~~~~^~~~~~~~~~~~~~
mergers.cpp:131:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
         scanf("%i",&colour[i]);
         ~~~~~^~~~~~~~~~~~~~~~~

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