Submission #846855

#	Time	Username	Problem	Language	Result	Execution time	Memory
846855		ogkostya	Beech Tree (IOI23_beechtree)	C++17	22 / 100	75 ms	13392 KiB

beechtree

#include "beechtree.h"
#include <algorithm> 
#include <queue>
#include <set>
#include <utility>


int cn[200010];
int lev[200010];
std::vector<int> down[200010];

std::vector<int> beechtree(int N, int M, std::vector<int> P, std::vector<int> C)
{
    std::vector<int> ans{  };

    bool sub2 = true;
    for (int i = 1; sub2 && i < N; i++)
    {
        sub2 &= P[i] == i - 1;
    }

    if (sub2)
    {
        ans.push_back(1);
        ans.push_back(1);
        for (int i = N - 3; i >= 0; i--)
        {
            sub2 &= C[i + 1] == C[N - 1];
            ans.push_back(sub2 ? 1 : 0);
        }
        std::reverse(ans.begin(), ans.end());
        return ans;
    }

    // sub 4
    std::fill_n(cn, M + 1, 0);
    bool sub4 = true;
    for (size_t i = 0; i < C.size(); i++)
    {
        if (++cn[C[i]] > 2)
        {
            sub4 = false;
            break;
        }
    }

    for (int i = 1; i < N; i++)
    {
        down[P[i]].push_back(i);
    }
    std::queue<int> q = {};
    std::fill_n(lev, N + 1, 0);
    std::fill_n(cn, N + 1, 0);
    for (int i = 0; i < N; i++)
    {
        cn[i] = down[i].size();
        if (down[i].size() == 0)
        {
            q.push(i);
        }
    }
    while (q.size() > 0)
    {
        int ii = q.front();
        q.pop();

        if (ii == 0)
            break;

        if (--cn[P[ii]] == 0)
        {
            q.push(P[ii]);
            lev[P[ii]] = 1 + lev[ii];
        }
    }

    for (int i = 0; i < N; i++)
    {
        if (i == 0 && !sub4 && lev[i] == 2)
        {
            std::set<int> s = {};
            bool ok = true;
            std::vector<std::pair<int, int>> sort = {};
            for (size_t j = 0; j < down[0].size(); j++)
            {
                if (s.count(C[down[0][j]]) == 0)
                {
                    s.insert(C[down[0][j]]);
                    sort.push_back(std::make_pair(down[down[0][j]].size(), down[0][j]));
                }
                else
                {
                    ok = false;
                    break;
                }
            }
            if (ok)
            {
                std::sort(sort.begin(), sort.end());
                std::reverse(sort.begin(), sort.end());
                for (size_t k = 0; ok && k < sort.size(); k++)
                {
                    std::set<int> next = {};
                    int index = sort[k].second;
                    for (size_t j = 0; j < down[index].size(); j++)
                    {
                        if (s.count(C[down[index][j]]) == 1 && next.count(C[down[index][j]]) == 0)
                        {
                            next.insert(C[down[index][j]]);
                        }
                        else
                        {
                            ok = false;
                            break;
                        }
                    }
                    s = next;
                }
            }
            ans.push_back(ok ? 1 : 0);
        }
        else if (lev[i] == 0)
        {
            ans.push_back(1);
        }
        else if (lev[i] == 1)
        {
            std::set<int> s = {};
            bool ok = true;
            for (size_t j = 0; j < down[i].size(); j++)
            {
                if (s.count(C[down[i][j]]) == 0)
                {
                    s.insert(C[down[i][j]]);
                }
                else
                {
                    ok = false;
                    break;
                }
            }
            ans.push_back(ok ? 1 : 0);
        }
        else if (lev[i] == 2)
        {
            std::set<int> s = {};
            bool ok = true;
            int f = -1;
            for (size_t j = 0; j < down[i].size(); j++)
            {
                if (down[down[i][j]].size() > 0)
                {
                    if (f == -1)
                    {
                        f = down[i][j];
                    }
                    else
                    {
                        ok = false;
                        break;
                    }
                }
                if (s.count(C[down[i][j]]) == 0)
                {
                    s.insert(C[down[i][j]]);
                }
                else
                {
                    ok = false;
                    break;
                }
            }
            if (ok)
            {
                for (size_t j = 0; j < down[f].size(); j++)
                {
                    if (s.count(C[down[f][j]]) == 1)
                    {
                        s.erase(C[down[f][j]]);
                    }
                    else
                    {
                        ok = false;
                        break;
                    }
                }
            }
            ans.push_back(ok ? 1 : 0);
        }
        else
        {
            ans.push_back(0);
        }
    }

    return ans;
}

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