Submission #945224

#	Time	Username	Problem	Language	Result	Execution time	Memory
945224		Nhoksocqt1	Beech Tree (IOI23_beechtree)	C++17	9 / 100	81 ms	40884 KiB

beechtree

#ifndef Nhoksocqt1
    #include "beechtree.h"
#endif // Nhoksocqt1
#include<bits/stdc++.h>
using namespace std;

#define inf 0x3f3f3f3f
#define sz(x) int((x).size())
#define fi first
#define se second
typedef long long ll;
typedef pair<int, int> ii;

template<class X, class Y>
	inline bool maximize(X &x, const Y &y) {return (x < y ? x = y, 1 : 0);}
template<class X, class Y>
	inline bool minimize(X &x, const Y &y) {return (x > y ? x = y, 1 : 0);}

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
int Random(int l, int r) {
    return uniform_int_distribution<int>(l, r)(rng);
}

const bool isMultiTest = 0;
const int MAXN = 200005;

vector<int> adj[MAXN], B[MAXN];
int maxDepth[MAXN];
int depth[MAXN], par[MAXN], col[MAXN], cnt[MAXN], lastQuery[MAXN], numNode, numColor;
bool canChoose[MAXN], allEqual[MAXN];

void preDfs(int u) {
    canChoose[u] = 1;
    for (int it = 0; it < sz(adj[u]); ++it) {
        int v(adj[u][it]);
        if(lastQuery[col[v]])
            canChoose[u] = 0;

        lastQuery[col[v]] = 1;
    }

    for (int it = 0; it < sz(adj[u]); ++it) {
        int v(adj[u][it]);
        lastQuery[col[v]] = 0;
    }

    allEqual[u] = 1;
    maxDepth[u] = depth[u];
    for (int it = 0; it < sz(adj[u]); ++it) {
        int v(adj[u][it]);
        depth[v] = depth[u] + 1;

        preDfs(v);
        maxDepth[u] = max(maxDepth[u], maxDepth[v]);
        canChoose[u] &= canChoose[v];
        allEqual[u] &= allEqual[v] & (sz(adj[v]) == 0 || col[v] == col[adj[v][0]]) & (col[v] == col[adj[u][0]]);
    }
}

void dfs1(int u) {
    B[u].clear();
    B[u].push_back(u);
    for (int it = 0; it < sz(adj[u]); ++it) {
        int v(adj[u][it]);
        dfs1(v);
        for (int jt = 0; jt < sz(B[v]); ++jt)
            B[u].push_back(B[v][jt]);
    }
}

vector<int> sub1(void) {
    dfs1(0);

    vector<int> ans;
    for (int i = 0; i < numNode; ++i) {
        sort(B[i].begin() + 1, B[i].end());
        bool check(0);
        do {
            for (int it = 0; it < sz(B[i]); ++it)
                cnt[col[B[i][it]]] = 0;

            check = 1;
            for (int it = 1; it < sz(B[i]); ++it) {
                int u(B[i][it]);
                if(B[i][cnt[col[u]]] != par[u]) {
                    check = 0;
                    break;
                }

                ++cnt[col[u]];
            }

            if(check)
                break;

        } while(next_permutation(B[i].begin() + 1, B[i].end()));
        ans.push_back(check);
    }

    return ans;
}

vector<int> subn2(void) {
    vector<int> ans;
    for (int i = 0; i < numNode; ++i) {
        if(!canChoose[i]) {
            ans.push_back(0);
            continue;
        }

        if(maxDepth[i] - depth[i] <= 1) {
            ans.push_back(1);
            continue;
        }

        if(maxDepth[i] - depth[i] > 2) {
            ans.push_back(allEqual[i]);
            continue;
        }

        for (int j = 0; j <= numNode; ++j) {
            B[j].clear();
            lastQuery[col[j]] = 0;
        }

        B[sz(adj[i])].push_back(i);

        int cnt(0);
        bool check(1);
        for (int j = numNode; j > 0; --j) {
            for (int it = 0; it < sz(B[j]); ++it) {
                int u(B[j][it]);
                for (int jt = 0; jt < sz(adj[u]); ++jt) {
                    int v(adj[u][jt]);
                    if(lastQuery[col[v]] != cnt || sz(adj[v]) > j) {
                        cout << '.' << sz(adj[v]) << ' ' << j << '\n';
                        check = 0;
                        break;
                    }

                    lastQuery[col[v]] = cnt + 1;
                    B[sz(adj[v])].push_back(v);
                }

                ++cnt;
                if(!check)
                    break;
            }

            if(!check)
                break;
        }

        ans.push_back(check);
    }

    //for (int it = 0; it < sz(ans); ++it) cout << ans[it]; cout << '\n';
    return ans;
}

vector<int> beechtree(int _N, int _M, vector<int> _P, vector<int> _C) {
    numNode = _N, numColor = _M;
    for (int i = 0; i < numNode; ++i) {
        par[i] = _P[i];
        col[i] = _C[i];
        if(i > 0)
            adj[par[i]].push_back(i);
    }

    preDfs(0);
    /*if(subn2() != sub1()) {
        cout << numNode << ' ' << numColor << '\n';
        for (int i = 0; i < numNode; ++i)
            cout << par[i] << ' '; cout << '\n';
        for (int i = 0; i < numNode; ++i)
            cout << col[i] << ' '; cout << '\n';

        exit(1);
    }*/

    //subn2();
    if(numNode <= 8) {
        return sub1();
    } else {
        return subn2();
    }
}

#ifdef Nhoksocqt1

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

    #define TASK "beechtree"
    if(fopen(TASK".inp", "r")) {
        freopen(TASK".inp", "r", stdin);
        freopen(TASK".out", "w", stdout);
    }

    vector<int> P, C;
    int N, M;
    cin >> N >> M;

    P.resize(N);
    C.resize(N);
    for (int i = 0; i < N; ++i) {
        cin >> P[i];
        //P[i] = (i == 0 ? -1 : Random(0, i - 1)); cout << P[i] << " \n"[i + 1 == N];
    }

    for (int i = 0; i < N; ++i) {
        cin >> C[i];
        //C[i] = Random(1, M); cout << C[i] << " \n"[i + 1 == N];
    }

    vector<int> ans = beechtree(N, M, P, C);
    cout << "ANSWER: ";
    for (int it = 0; it < sz(ans); ++it)
        cout << ans[it]; cout << '\n';

    return 0;
}

#endif // Nhoksocqt1

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