Submission #1095368

#	Time	Username	Problem	Language	Result	Execution time	Memory
1095368		raphaelp	Team Coding (EGOI24_teamcoding)	C++14	100 / 100	2341 ms	62672 KiB

Main

#include <bits/stdc++.h>
using namespace std;
void dfs(long long x, vector<vector<long long>> &AR, vector<long long> &depth, long long cur)
{
    depth[x] = cur;
    for (long long i = 0; i < AR[x].size(); i++)
    {
        dfs(AR[x][i], AR, depth, cur + 1);
    }
}
void dfs2(long long x, vector<vector<long long>> &AR, vector<vector<pair<long long, long long>>> &colorpos, vector<long long> &available, map<long long, long long> &M, vector<long long> &C, vector<vector<long long>> &colors, long long cur, long long &best, long long &best2)
{
    available.push_back(1);
    vector<long long> available2;
    map<long long, long long> M2;
    M[C[x]] += 1;
    for (long long i = 0; i < AR[x].size(); i++)
    {
        available2.clear();
        M2.clear();
        dfs2(AR[x][i], AR, colorpos, available2, M2, C, colors, cur + 1, best, best2);
        available2.push_back(0);
        if (available2.size() > available.size())
            swap(available2, available);
        for (long long j = 0; j < available2.size(); j++)
        {
            available[j + available.size() - available2.size()] += available2[j];
        }
        if (M2.size() > M.size())
            swap(M, M2);
        for (auto j : M2)
        {
            M[j.first] += j.second;
        }
    }
    if (colors[C[x]].size() <= sqrt(AR.size()))
    {
        long long ans1 = 0, ans2 = 0;
        for (long long i = 0; i < colorpos[C[x]].size(); i++)
        {
            if (colorpos[C[x]][i].first >= cur && colorpos[C[x]][i].first < cur + available.size())
                ans1 += min(colorpos[C[x]][i].second, available[available.size() - colorpos[C[x]][i].first - 1 + cur]);
        }
        ans2 = ans1 - M[C[x]];
        if (ans1 > best || (ans1 == best && ans2 < best2))
        {
            best = ans1, best2 = ans2;
        }
    }
}
pair<long long, long long> dfs3(long long x, vector<vector<long long>> &AR, vector<long long> &C, long long c, vector<long long> &used, vector<long long> &maxx, long long cur, long long &best, long long &best2)
{
    vector<long long> used2;
    long long tot = 0, tot2 = 0;
    if (C[x] == c)
        tot2++;
    used.push_back(0);
    if (maxx[cur] != 0)
    {
        used[0] = 1;
        tot++;
    }
    for (long long i = 0; i < AR[x].size(); i++)
    {
        used2.clear();
        pair<long long, long long> temp = dfs3(AR[x][i], AR, C, c, used2, maxx, cur + 1, best, best2);
        tot2 += temp.second;
        used2.push_back(0);
        if (used2.size() > used.size())
        {
            swap(used2, used);
            swap(tot, temp.first);
        }
        for (long long j = 0; j < used2.size(); j++)
        {
            tot -= used[j + used.size() - used2.size()];
            used[j + used.size() - used2.size()] = min(used[j + used.size() - used2.size()] + used2[j], maxx[cur + used2.size() - j - 1]);
            tot += used[j + used.size() - used2.size()];
        }
    }
    if (C[x] == c && (tot > best || (tot == best && tot - tot2 < best2)))
    {
        best = tot, best2 = tot - tot2;
    }
    return {tot, tot2};
}
int main()
{
    cin.tie(0);
    cout.tie(0);
    ios_base::sync_with_stdio(0);
    long long N, K;
    cin >> N >> K;
    vector<long long> C(N);
    vector<vector<long long>> colors(K);
    for (long long i = 0; i < N; i++)
    {
        cin >> C[i];
        colors[C[i]].push_back(i);
    }
    vector<long long> p(N, 0);
    vector<vector<long long>> AR(N);
    for (long long i = 1; i < N; i++)
    {
        cin >> p[i];
        AR[p[i]].push_back(i);
    }
    vector<long long> depth(N);
    vector<vector<pair<long long, long long>>> colorpos(K);
    dfs(0, AR, depth, 0);
    long long best = 0, best2 = 0;
    for (long long i = 0; i < K; i++)
    {
        for (long long j = 0; j < colors[i].size(); j++)
        {
            colorpos[i].push_back({depth[colors[i][j]], 1});
        }
        sort(colorpos[i].begin(), colorpos[i].end());
        vector<pair<long long, long long>> temp;
        long long cur = 0;
        for (long long j = 0; j < colors[i].size(); j++)
        {
            if (j > 0 && colorpos[i][j].first != colorpos[i][j - 1].first)
            {
                temp.push_back({colorpos[i][j - 1].first, cur});
                cur = 0;
            }
            cur++;
        }
        if (colors[i].size() > 0)
            temp.push_back({colorpos[i].back().first, cur});
        swap(temp, colorpos[i]);
    }
    vector<long long> available;
    map<long long, long long> M;
    dfs2(0, AR, colorpos, available, M, C, colors, 0, best, best2);
    for (long long i = 0; i < K; i++)
    {
        if (colors[i].size() >= sqrt(N) - 1)
        {
            vector<long long> used, maxx(N);
            for (long long j = 0; j < colors[i].size(); j++)
                maxx[depth[colors[i][j]]]++;
            pair<long long, long long> temp = dfs3(0, AR, C, i, used, maxx, 0, best, best2);
        }
    }
    cout << best << ' ' << best2;
}

Compilation message (stderr)

Main.cpp: In function 'void dfs(long long int, std::vector<std::vector<long long int> >&, std::vector<long long int>&, long long int)':
Main.cpp:6:29: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
    6 |     for (long long i = 0; i < AR[x].size(); i++)
      |                           ~~^~~~~~~~~~~~~~
Main.cpp: In function 'void dfs2(long long int, std::vector<std::vector<long long int> >&, std::vector<std::vector<std::pair<long long int, long long int> > >&, std::vector<long long int>&, std::map<long long int, long long int>&, std::vector<long long int>&, std::vector<std::vector<long long int> >&, long long int, long long int&, long long int&)':
Main.cpp:17:29: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   17 |     for (long long i = 0; i < AR[x].size(); i++)
      |                           ~~^~~~~~~~~~~~~~
Main.cpp:25:33: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   25 |         for (long long j = 0; j < available2.size(); j++)
      |                               ~~^~~~~~~~~~~~~~~~~~~
Main.cpp:39:33: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   39 |         for (long long i = 0; i < colorpos[C[x]].size(); i++)
      |                               ~~^~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:41:75: warning: comparison of integer expressions of different signedness: 'long long int' and 'long long unsigned int' [-Wsign-compare]
   41 |             if (colorpos[C[x]][i].first >= cur && colorpos[C[x]][i].first < cur + available.size())
Main.cpp: In function 'std::pair<long long int, long long int> dfs3(long long int, std::vector<std::vector<long long int> >&, std::vector<long long int>&, long long int, std::vector<long long int>&, std::vector<long long int>&, long long int, long long int&, long long int&)':
Main.cpp:63:29: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   63 |     for (long long i = 0; i < AR[x].size(); i++)
      |                           ~~^~~~~~~~~~~~~~
Main.cpp:74:33: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   74 |         for (long long j = 0; j < used2.size(); j++)
      |                               ~~^~~~~~~~~~~~~~
Main.cpp: In function 'int main()':
Main.cpp:114:33: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  114 |         for (long long j = 0; j < colors[i].size(); j++)
      |                               ~~^~~~~~~~~~~~~~~~~~
Main.cpp:121:33: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  121 |         for (long long j = 0; j < colors[i].size(); j++)
      |                               ~~^~~~~~~~~~~~~~~~~~
Main.cpp:142:37: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  142 |             for (long long j = 0; j < colors[i].size(); j++)
      |                                   ~~^~~~~~~~~~~~~~~~~~
Main.cpp:144:40: warning: variable 'temp' set but not used [-Wunused-but-set-variable]
  144 |             pair<long long, long long> temp = dfs3(0, AR, C, i, used, maxx, 0, best, best2);
      |                                        ^~~~

• Subtask 1: 12 / 12
• Subtask 2: 19 / 19
• Subtask 3: 27 / 27
• Subtask 4: 23 / 23
• Subtask 5: 19 / 19