Submission #848079

#	Time	Username	Problem	Language	Result	Execution time	Memory
848079		ogkostya	Beech Tree (IOI23_beechtree)	C++17	31 / 100	103 ms	18884 KiB

beechtree

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

#define MAX 200030

int cn[MAX];
int lev[MAX];
int c2type[MAX];
int c2count0[MAX];
int c2count1[MAX];
int c2tail0[MAX];
int c2tail1[MAX];
int c2ans[MAX];
std::vector<int> down[MAX];

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;
    }

    std::fill_n(cn, M + 1, 0);
    for (size_t i = 0; i < C.size(); i++)
    {
        cn[C[i]]++;
    }
    std::vector<int> cc{  };
    for (size_t i = 1; i <= M; i++)
    {
        if (cn[i] > 0)
        {
            cc.push_back(i);
        }
    }
    for (size_t i = 0; i < C.size(); i++)
    {
        std::vector<int>::iterator ind = std::lower_bound(cc.begin(), cc.end(), C[i]);
        C[i] = ind - cc.begin();
    }

    bool sub4 = true;
    for (size_t i = 1; i <= M; i++)
    {
        if (cn[i] > 2)
        {
            sub4 = false;
            break;
        }
    }

    for (int i = 0; i < N; i++)
    {
        down[i].clear();
    }

    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(c2count0, N + 1, 0);
    std::fill_n(c2count1, N + 1, 0);
    std::fill_n(c2tail0, N + 1, 0);
    std::fill_n(c2tail1, N + 1, 0);
    std::fill_n(c2type, N + 1, 0);
    std::fill_n(c2ans, 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);
        }
    }

    int cc0 = 0;
    int cc1 = 0;
    int cc2 = 0;
    int cc3 = 0;

    bool c2 = cc.size() == 2;
    while (q.size() > 0)
    {
        int ii = q.front();
        q.pop();

        if (c2)
        {
            if (down[ii].size() == 0)
            {
                cc0++;
                c2ans[ii] = 1;
            }
            else if (down[ii].size() == 1)
            {
                cc1++;
                int t = down[ii][0];
                c2ans[ii] = c2ans[t];
                if (c2ans[t] == 1)
                {
                    if (C[t] == 0)
                    {
                        if (c2count1[t] != 0)
                        {
                            c2ans[ii] = 0;
                        }
                        else
                        {
                            c2count0[ii] = 1 + c2count0[t];
                            c2tail0[ii] = c2count0[ii];
                        }
                    }
                    else
                    {
                        if (c2count0[t] != 0)
                        {
                            c2ans[ii] = 0;
                        }
                        else
                        {
                            c2count1[ii] = 1 + c2count1[t];
                            c2tail1[ii] = c2count1[ii];
                        }
                    }
                }
            }
            else if (down[ii].size() == 2)
            {
                cc2++;
                int i0 = down[ii][0];
                int i1 = down[ii][1];
                if (C[i0] + C[i1] != 1 || c2ans[i0] + c2ans[i1] != 2)
                {
                    c2ans[ii] = 0;
                }
                else
                {
                    if (C[i0] == 1)
                    {
                        std::reverse(down[ii].begin(), down[ii].end());
                        int xxx = i0;
                        i0 = i1;
                        i1 = xxx;
                    }
                    int i01 = -1;
                    for (size_t i = 0; i < down[i0].size(); i++)
                    {
                        if (C[down[i0][i]] == 1)
                            i01 = down[i0][i];
                    }
                    int i10 = -1;
                    for (size_t i = 0; i < down[i1].size(); i++)
                    {
                        if (C[down[i1][i]] == 0)
                            i10 = down[i1][i];
                    }
                    if (lev[ii] == 1)
                    {
                        c2count0[ii] = 1;
                        c2count1[ii] = 1;
                        c2ans[ii] = 1;
                    }
                    else if (c2tail0[i1] > 1 && (c2count0[i0] - c2count1[i0] < 0 || c2count0[i0] - c2count1[i0] == 0 && c2count0[i0] != 0))
                    {
                        c2ans[ii] = 0;
                    }
                    else if (c2tail0[i0] > 1 && (c2count0[i1] - c2count1[i1] < 0 || c2count0[i1] - c2count1[i1] == 0 && c2count0[i1] != 0))
                    {
                        c2ans[ii] = 0;
                    }
                    else if (c2tail1[i0] > 1 && (c2count1[i1] - c2count0[i1] < 0 || c2count1[i1] - c2count0[i1] == 0 && c2count1[i1] != 0))
                    {
                        c2ans[ii] = 0;
                    }
                    else if (c2tail1[i1] > 1 && (c2count1[i0] - c2count0[i0] < 0 || c2count1[i0] - c2count0[i0] == 0 && c2count1[i0] != 0))
                    {
                        c2ans[ii] = 0;
                    }
                    else if (c2count0[i1] - c2count1[i1] > 0 && c2count1[i0] - c2count0[i0] > 0)
                    {
                        c2ans[ii] = 0;
                    }
                    else if (c2count1[i1] - c2count0[i1] > 0 && c2count0[i0] - c2count1[i0] > 0)
                    {
                        c2ans[ii] = 0;
                    }
                    else if (c2count0[i1] - c2count1[i1] > 0 && c2count1[i0] - c2count0[i0] > 0)
                    {
                        c2ans[ii] = 0;
                    }
                    else if (i01 != -1 && (c2count0[i1] < c2count0[i01] || c2count1[i1] < c2count1[i01]))
                    {
                        c2ans[ii] = 0;
                    }
                    else if (i10 != -1 && (c2count0[i0] < c2count0[i10] || c2count1[i0] < c2count1[i10]))
                    {
                        c2ans[ii] = 0;
                    }
                    else
                    {
                        c2ans[ii] = 1;
                        c2count0[ii] = 1 + c2count0[i0] + c2count0[i1];
                        c2count1[ii] = 1 + c2count1[i0] + c2count1[i1];
                        c2tail0[ii] = std::max(c2tail0[i0], c2tail0[i1]);
                        c2tail1[ii] = std::max(c2tail1[i0], c2tail1[i1]);
                    }
                }
            }
            else
            {
                cc3++;
                c2ans[ii] = 0;
            }
        }

        if (ii == 0)
            break;

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

    if (c2)
    {
        for (int i = 0; i < N; i++)
        {
            ans.push_back(c2ans[i]);
        }
        //printf("%d %d %d %d\n", cc0, cc1, cc2, cc3);
    }
    /*if (N <= 8)
    {

    }*/
    else
    {
        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;
}

Compilation message (stderr)

beechtree.cpp: In function 'std::vector<int> beechtree(int, int, std::vector<int>, std::vector<int>)':
beechtree.cpp:48:26: warning: comparison of integer expressions of different signedness: 'size_t' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   48 |     for (size_t i = 1; i <= M; i++)
      |                        ~~^~~~
beechtree.cpp:62:26: warning: comparison of integer expressions of different signedness: 'size_t' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   62 |     for (size_t i = 1; i <= M; i++)
      |                        ~~^~~~
beechtree.cpp:185:118: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  185 |                     else if (c2tail0[i1] > 1 && (c2count0[i0] - c2count1[i0] < 0 || c2count0[i0] - c2count1[i0] == 0 && c2count0[i0] != 0))
      |                                                                                     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~
beechtree.cpp:189:118: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  189 |                     else if (c2tail0[i0] > 1 && (c2count0[i1] - c2count1[i1] < 0 || c2count0[i1] - c2count1[i1] == 0 && c2count0[i1] != 0))
      |                                                                                     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~
beechtree.cpp:193:118: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  193 |                     else if (c2tail1[i0] > 1 && (c2count1[i1] - c2count0[i1] < 0 || c2count1[i1] - c2count0[i1] == 0 && c2count1[i1] != 0))
      |                                                                                     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~
beechtree.cpp:197:118: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  197 |                     else if (c2tail1[i1] > 1 && (c2count1[i0] - c2count0[i0] < 0 || c2count1[i0] - c2count0[i0] == 0 && c2count1[i0] != 0))
      |                                                                                     ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~

• Subtask 0: 0 / 0
• Subtask 1: 9 / 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