Submission #1233016

#	Time	Username	Problem	Language	Result	Execution time	Memory
1233016		antonn	Beech Tree (IOI23_beechtree)	C++20	49 / 100	2098 ms	45640 KiB

beechtree

#include "beechtree.h"
#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

const int N = 2e5 + 7;
const int L = 20;

int n, m;
vector<int> g[N];
int col[N], ok[N], par[N], bad[N], sz[N];
int anc[L][N], dep[N];

void build(int u) {
    sz[u] = 1;
    for (auto v : g[u]) {
        dep[v] = dep[u] + 1;
        build(v);
        sz[u] += sz[v];
    }
}

int lca(int a, int b) {
    if (dep[a] < dep[b]) swap(a, b);
    for (int h = L - 1; h >= 0; --h) {
        if (dep[a] - (1 << h) >= dep[b]) {
            a = anc[h][a];
        }
    }
    if (a == b) return a;
    for (int h = L - 1; h >= 0; --h) {
        if (anc[h][a] != anc[h][b]) {
            a = anc[h][a];
            b = anc[h][b];
        }
    }
    return anc[0][a];
}

void dfs(int u) {
    for (auto v : g[u]) {
        dfs(v);
        bad[u] += bad[v];
    }
    if (g[u].size() == 0) {
        ok[u] = 1;
        return;
    }
    ok[u] = 1;
    if (bad[u]) ok[u] = 0;
    set<int> children;
    for (auto v : g[u]) children.insert(col[v]);
    if (children.size() != g[u].size()) ok[u] = 0;
    for (auto v : g[u]) ok[u] &= ok[v];
    
}

vector<int> beechtree(int N, int M, vector<int> P, vector<int> C) {
    n = N;
    m = M;
    for (int i = 0; i < n; ++i) {
        par[i] = P[i];
        if (i > 0) {
            g[P[i]].push_back(i);
        }
    }
    for (int i = 0; i < n; ++i) col[i] = C[i];
    
    for (int i = 0; i < n; ++i) {
        anc[0][i] = par[i];
    }
    for (int h = 1; h < L; ++h) {
        for (int i = 0; i < n; ++i) {
            if (anc[h - 1][i] == -1) {
                anc[h][i] = -1;
            } else {
                anc[h][i] = anc[h - 1][anc[h - 1][i]];
            }
        }
    }
    
    build(0);
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < n; ++j) {
            if (sz[i] >= sz[j]) {
                map<int, int> A, B;
                for (auto x : g[i]) A[col[x]] = sz[x];
                for (auto x : g[j]) B[col[x]] = sz[x];
                bool ok = 1;
                for (auto it : B) ok &= (it.second <= A[it.first]);
                if (!ok) {
                    bad[lca(i, j)]++;
                }
            }
        }
    }
    
    dfs(0);
    vector<int> ret(n);
    for (int i = 0; i < n; ++i) ret[i] = ok[i];
    return ret; 
}

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