Submission #1156593

#	Time	Username	Problem	Language	Result	Execution time	Memory
1156593		beaconmc	Team Coding (EGOI24_teamcoding)	C++20	100 / 100	1880 ms	41304 KiB

Main

// not my code - I write better code than this

#include <bits/stdc++.h>

using namespace std;

const int X = 500;

int main() {
  cin.tie(0)->sync_with_stdio(0);
  int N, K;
  cin >> N >> K;
  /*
    for each node u, the answer is:
    - let the number of people on the i-th layer in the tree liking language l be b_li
    - let the number of people on the i-th layer in the subtree of u be a_ui
    - let the number of people in the subtree of u liking language L[u] be c_u
    - ans[u][0] = sum of min(a_ui, b_li) for all i
    - ans[u][1] = ans[0] - c_u
  */
  /*
    for each language:
    - if number who prefer >= X:
      -> at most N/X of these
      -> run O(N log N) small to large algorithm
      -> algorithm O(N log N * (N/X))
    - if number who prefer < X
      -> run O(X^2 log N) algorithm
      -> so altogether is O(N/X * X^2 log N) = O(NX log N)
    - so want N*N*logN/X = N*X*logN => N=X*X => X=sqrt(N) = 320
    - now NX*logN ~ 5 * 10^8, so maybe
  */
  vector<int> L(N), LC(K);
  vector<vector<int>> LU(K);
  for (int i=0; i<N; i++) {
    cin >> L[i];
    LC[L[i]] += 1;
    LU[L[i]].push_back(i);
  }
  vector<vector<int>> adj(N);
  for (int i=0; i<N-1; i++) {
    int u = i + 1, v;
    cin >> v;
    adj[u].push_back(v);
    adj[v].push_back(u);
  }
  /*
    precalculate:
    - in[u], out[u], depth[u], b[l]
    - vector of ins for each node at depth level
    - to count number of nodes at depth d in subtree of u:
      -> count number between in[u] and out[u] in the in-depth vector
    - 
  */
  vector<int> depth(N), in(N), out(N);
  vector<map<int,int>> B(K);
  vector<vector<int>> ins(N);
  int T = -1;
  function<void(int,int)> dfs = [&](int u, int p) -> void {
    in[u] = ++T;
    B[L[u]][depth[u]] += 1;
    ins[depth[u]].push_back(in[u]);
    for (int v : adj[u]) {
      if (v == p) continue;
      depth[v] = depth[u] + 1;
      dfs(v, u);
    }
    out[u] = T;
  };
  dfs(0, 0);
  function<int(int,int)> count = [&](int u, int depth) -> int {
    vector<int> &vec = ins[depth];
    auto itL = lower_bound(vec.begin(), vec.end(), in[u]);
    auto itR = upper_bound(vec.begin(), vec.end(), out[u]);
    return itR - itL;
  };
  array<int,2> ans = {1, 1};
  for (int i=0; i<K; i++) {
    if (LC[i] < X) {
      // run O(X^2 log N) algorithm
      // for each node, count:
      // -> the number of nodes on each depth in a subtree where 
      //    another of the X nodes is at the same depth - O(X)
      // -> the number of nodes inside subtree already
      // -> the number of the X nodes at each depth
      for (int u : LU[i]) {
        map<int,int> A;
        int res = 1, c = 1;
        for (int v : LU[i]) {
          if (depth[v] > depth[u]) {
            A[depth[v]] += 1;
            if (in[u] <= in[v] && in[v] <= out[u]) {
              c += 1;
            }
          }
        }
        for (auto &[dep, a] : A) {
          res += min(a, count(u, dep));
        }
        ans = max(ans, {res, c});
      }
    } else {
      // run O(N log N) algorithm
      // for each node, count:
      // -> the number of nodes on each depth in a subtree - small-to-large map
      // -> but, limit the above counter to the number of nodes of language at that depth
      // -> the number of nodes of the same language in a subtree - trivial
      vector<map<int,int>> M(N);
      vector<array<int,2>> res(N);
      function<void(int,int)> merge = [&](int u, int v) -> void {
        if (M[u].size() < M[v].size()) swap(M[u], M[v]), swap(res[u], res[v]);
        for (auto &[dep, ct] : M[v]) {
          int next = min(ct + M[u][dep], B[i][dep]);
          res[u][0] += next - M[u][dep];
          M[u][dep] = next;
        }
        res[u][1] += res[v][1];
      };
      function<void(int,int)> dfs = [&](int u, int p) -> void {
        for (int v : adj[u]) {
          if (v == p) continue;
          dfs(v, u);
          merge(u, v);
        }
        if (L[u] == i) {
          res[u][0] += 1;
          res[u][1] += 1;
          M[u][depth[u]] = 1;
          ans = max(ans, res[u]);
        } else if (B[i][depth[u]]) {
          res[u][0] += 1;
          M[u][depth[u]] = 1;
        }
      };
      dfs(0, 0);
    }
  }
  cout << ans[0] << " " << ans[0]-ans[1] << "\n";
  return 0;
}

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