Submission #952844

#	Time	Username	Problem	Language	Result	Execution time	Memory
952844		kilkuwu	Regions (IOI09_regions)	C++17	100 / 100	2894 ms	127224 KiB

regions

#include <bits/stdc++.h>


#define nl '\n'

#ifdef LOCAL
#include "template\debug.hpp"
#else
#define dbg(...) ;
#define timer(...) ;
#endif

signed main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(nullptr);

  int n, r, q;
  std::cin >> n >> r >> q;
  std::vector<int> home(n), pa(n);
  std::vector<std::vector<int>> adj(n);
  std::vector<std::vector<int>> homes(r);
  std::cin >> home[0];
  --home[0];
  homes[home[0]].push_back(0);
  for (int i = 1; i < n; i++) {
    std::cin >> pa[i] >> home[i];
    adj[--pa[i]].push_back(i);
    --home[i];
    homes[home[i]].push_back(i);
  }

  std::vector<int> st(n), en(n);
  int ti = 0;

  std::function<void(int)> dfs = [&](int u) -> void {
    st[u] = ti++;
    for (int v : adj[u]) {
      dfs(v);
    }
    en[u] = ti;
  };


  dfs(0);

  for (int i = 0; i < r; i++) {
    std::sort(homes[i].begin(), homes[i].end(), [&](int u, int v) {
      return st[u] < st[v];
    });
  }

  constexpr int BLOCK_SIZE = 400;

  std::vector<std::vector<int>> b1((n - 1) / BLOCK_SIZE + 1, std::vector<int>(r));
  std::vector<std::vector<int>> b2((n - 1) / BLOCK_SIZE + 1, std::vector<int>(r));
  std::vector<int> map_block(r, -1);
  std::vector<int> dp(n);

  int bid = 0;
  for (int x = 0; x < r; x++) {
    if (homes[x].size() >= BLOCK_SIZE) {
      map_block[x] = bid;

      std::fill(dp.begin(), dp.end(), 0);

      for (int u : homes[x]) dp[u] = 1;

      for (int i = 0; i < n; i++) {
        if (i > 0) dp[i] += dp[pa[i]]; // dp[i] = number of us that is parent of this
        b1[bid][home[i]] += dp[i];
      }

      std::fill(dp.begin(), dp.end(), 0);

      for (int u : homes[x]) dp[u] = 1;

      for (int i = n - 1; i >= 0; i--) {
        if (i > 0) dp[pa[i]] += dp[i]; // dp[i] = number of us that this is parent
        b2[bid][home[i]] += dp[i]; // con den cha
      }
      
      bid++;
    }
  }

  while (q--) {
    int a, b;
    std::cin >> a >> b;
    --a, --b;

    if (map_block[a] != -1) {
      // how can i do the reverse?
      std::cout << b1[map_block[a]][b] << std::endl;
      continue;
    } 

    if (map_block[b] != -1) {
      std::cout << b2[map_block[b]][a] << std::endl;
      continue;
    }

    int ans = 0;
    for (int u : homes[a]) {
      int f = std::lower_bound(homes[b].begin(), homes[b].end(), st[u], [&](int x, int y) {
        return st[x] < y;
      }) - homes[b].begin();
      int e = std::lower_bound(homes[b].begin(), homes[b].end(), en[u], [&](int x, int y) {
        return st[x] < y;
      }) - homes[b].begin();

      ans += e - f;
    }

    std::cout << ans << std::endl;
  }
}      

/*

homes[1] = {1, 6}
homes[2] = {2}
homes[3] = {3, 4, 5}

*/

• Subtask 1: 15 / 15
• Subtask 2: 85 / 85