Submission #967124

#	Time	Username	Problem	Language	Result	Execution time	Memory
967124		sleepntsheep	Capital City (JOI20_capital_city)	C11	100 / 100	405 ms	24300 KiB

capital_city

#include <stdio.h>

#define N 200000
#define M (2*N-2)
#define N_ N+1
#define M_ M+1
#define K N

int n, k, color[N], answer;
/* graph */
int head[N], nxt[M_], vv[M_];

/* node with color */
int head_[K], nxt_[N_], vv_[N_];

void link(int*head, int*nxt, int*vv, int*i, int u, int v)
{
    nxt[i[0]] = head[u];
    vv[i[0]] = v;
    head[u] = i[0]++;
}

int dead[N], sz[N];

void dfs2(int u, int p)
{
    sz[u] = 1;
    for (int j = head[u]; j; j = nxt[j])
    {
        if (vv[j] == p || dead[vv[j]])
            continue;
        dfs2(vv[j], u);
        sz[u] += sz[vv[j]];
    }
}

int par[N], visit[K], dfreq[K], cfreq[K];

void dfs3(int u, int p)
{
    visit[color[u]] = -1;
    ++dfreq[color[u]];
    par[u] = p;
    for (int j = head[u]; j; j = nxt[j])
    {
        if (vv[j] == p || dead[vv[j]])
            continue;
        dfs3(vv[j], u);
    }
}

void dfs4(int u, int p)
{
    visit[color[u]] = -1;
    dfreq[color[u]] = 0;
    for (int j = head[u]; j; j = nxt[j])
    {
        if (vv[j] == p || dead[vv[j]])
            continue;
        dfs4(vv[j], u);
    }
}

int get_centroid(int u, int p, int tsize)
{
    for (int j = head[u]; j; j = nxt[j])
        if (!dead[vv[j]] && vv[j] != p && sz[vv[j]] * 2 >= tsize)
            return get_centroid(vv[j], u, tsize);
    return u;
}

void centroid_decomposition(int u)
{
    dfs2(u, u);
    u = get_centroid(u, u, sz[u]);
    dead[u] = 1;

    dfs3(u, u);

    int col = color[u];
    if (dfreq[col] == cfreq[col])
    {
        static int qu[N], head, tail;
        head = tail = 0;

        int dcolor = 0;
        for (int at, jj = head_[col]; jj; jj = nxt_[jj])
            qu[head++] = vv_[jj];
        visit[col] = 1;

        while (head-tail)
        {
            int uu = qu[tail++];
            if (uu == u || visit[color[par[uu]]] == 1)
                continue;
            int nextcol = color[par[uu]];
            if (cfreq[nextcol] != dfreq[nextcol])
                goto NOTOK;
            ++dcolor;
            for (int jj = head_[nextcol]; jj; jj = nxt_[jj])
                qu[head++] = vv_[jj];
            visit[nextcol] = 1;
        }

        if (dcolor < answer)
            answer = dcolor;
NOTOK:;
    }
    dfs4(u, u);

    for (int j = head[u]; j; j = nxt[j])
        if (!dead[vv[j]])
            centroid_decomposition(vv[j]);
}

int main()
{
    scanf("%d%d", &n, &k);
    answer = n;
    for (int i = 1, u, v, eo[1] = { 1 }; i < n; ++i)
    {
        scanf("%d%d", &u, &v); --u, --v;
        link(head, nxt, vv, eo, u, v);
        link(head, nxt, vv, eo, v, u);
    }
    for (int eo[1] = { 1 }, i = 0; i < n; ++i)
    {
        scanf("%d", color+i);
        ++cfreq[--color[i]];
        link(head_, nxt_, vv_, eo, color[i], i);
    }

    centroid_decomposition(0);

    printf("%d", answer);
}

Compilation message (stderr)

capital_city.c: In function 'centroid_decomposition':
capital_city.c:87:18: warning: unused variable 'at' [-Wunused-variable]
   87 |         for (int at, jj = head_[col]; jj; jj = nxt_[jj])
      |                  ^~
capital_city.c: In function 'main':
capital_city.c:118:5: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
  118 |     scanf("%d%d", &n, &k);
      |     ^~~~~~~~~~~~~~~~~~~~~
capital_city.c:122:9: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
  122 |         scanf("%d%d", &u, &v); --u, --v;
      |         ^~~~~~~~~~~~~~~~~~~~~
capital_city.c:128:9: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
  128 |         scanf("%d", color+i);
      |         ^~~~~~~~~~~~~~~~~~~~

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 30 / 30
• Subtask 4: 59 / 59