Submission #841420

#	Time	Username	Problem	Language	Result	Execution time	Memory
841420		model_code	Beech Tree (IOI23_beechtree)	C++17	100 / 100	224 ms	35696 KiB

beechtree

// correct/solution-3.cpp

#include "beechtree.h"

#include <map>
#include <set>

using namespace std;
using pii = pair<int, int>;

int n, m;
vector<int> par, col, sz, ord;
vector<vector<int>> ch;
vector<bool> rem;
vector<int> sz_rem;
set<pair<pii, int>> q;
map<int, int> cnt;
vector<int> t;

bool chk(int u)
{
    ord.clear();
    q.insert({{-sz[u], 0}, u});
    while (!q.empty())
    {
        int x = q.begin()->second;
        q.erase(q.begin());
        ord.push_back(x);
        for (int y : ch[x])
        {
            q.insert({{-sz[y], (int)ord.size()}, y});
        }
    }
    cnt.clear();
    for (int k = 1; k < (int)ord.size(); ++k)
    {
        int v = ord[k];
        int l = (k == 1 ? 0 : cnt[col[v]]);
        if (par[v] != ord[l])
        {
            return false;
        }
        cnt[col[v]]++;
    }
    return true;
}

int dfs(int u, int p = -1)
{
    sz_rem[u] = 1;
    for (int v : ch[u])
    {
        if (v != p && !rem[v])
        {
            sz_rem[u] += dfs(v, u);
        }
    }
    int v = par[u];
    if (v >= 0 && v != p && !rem[v])
    {
        sz_rem[u] += dfs(v, u);
    }
    return sz_rem[u];
}

int get_centroid(int u, int csz, int p = -1)
{
    for (int v : ch[u])
    {
        if (v != p && !rem[v])
        {
            if (sz_rem[v] * 2 > csz)
            {
                return get_centroid(v, csz, u);
            }
        }
    }
    int v = par[u];
    if (v >= 0 && v != p && !rem[v])
    {
        if (sz_rem[v] * 2 > csz)
        {
            return get_centroid(v, csz, u);
        }
    }
    return u;
}

void solve(int r)
{
    int c = get_centroid(r, dfs(r));

    if (t[c] == -1)
    {
        t[c] = chk(c);
        if (t[c] == 0)
            for (int v = par[c]; v >= 0 && t[v] == -1; v = par[v])
            {
                t[v] = 0;
            }
    }

    rem[c] = true;
    if (t[c] != 1)
    {
        for (int v : ch[c])
        {
            if (!rem[v])
            {
                solve(v);
            }
        }
    }
    int v = par[c];
    if (v >= 0 && !rem[v])
        solve(v);
}

std::vector<int> beechtree(int N, int M, std::vector<int> P, std::vector<int> C)
{
    n = N, m = M, par = P, col = C;
    ch.assign(N, {});
    sz.assign(N, 1);
    t.assign(N, -1);
    for (int v = 1; v < N; ++v)
    {
        int u = par[v];
        ch[u].push_back(v);
    }
    for (int v = N - 1; v > 0; --v)
    {
        sz[par[v]] += sz[v];
    }

    rem.assign(N, false), sz_rem.assign(N, 0);
    solve(0);
    for (int v = 1; v < N; ++v)
    {
        if (t[par[v]] == 1)
            t[v] = 1;
    }

    return t;
}

• Subtask 0: 0 / 0
• Subtask 1: 9 / 9
• Subtask 2: 5 / 5
• Subtask 3: 9 / 9
• Subtask 4: 8 / 8
• Subtask 5: 14 / 14
• Subtask 6: 14 / 14
• Subtask 7: 12 / 12
• Subtask 8: 17 / 17
• Subtask 9: 12 / 12