Submission #799712

#	Time	Username	Problem	Language	Result	Execution time	Memory
799712		tch1cherin	Cat in a tree (BOI17_catinatree)	C++17	100 / 100	234 ms	186872 KiB

catinatree

#include <bits/stdc++.h>
using namespace std;

struct segment_tree {
  int lx, rx, value = 0;
  segment_tree* left = nullptr;
  segment_tree* right = nullptr;

  segment_tree(int _lx, int _rx) : lx(_lx), rx(_rx) {}

  void extend_left() {
    if (!left) left = new segment_tree(lx, (lx + rx) / 2);
  }

  void extend_right() {
    if (!right) right = new segment_tree((lx + rx) / 2, rx);
  }

  void modify(int i, int v) {
    if (rx - lx == 1) {
      value = max(value, v);
    } else {
      int mid = (lx + rx) / 2;
      if (i < mid) {
        extend_left();
        left->modify(i, v);
      } else {
        extend_right();
        right->modify(i, v);
      }
      value = max(!left ? 0 : left->value, !right ? 0 : right->value);
    }
  }

  int get(int i) {
    if (rx - lx == 1) {
      return value;
    } else {
      int mid = (lx + rx) / 2;
      if (i < mid) {
        return !left ? 0 : left->get(i);
      } else {
        return !right ? 0 : right->get(i);
      }
    }
  }

  int query(int l) {
    if (rx <= l) {
      return 0;
    } else if (lx >= l) {
      return value;
    } else {
      return max(!left ? 0 : left->query(l), !right ? 0 : right->query(l));
    }
  }
};

const int MAX_N = 2e5 + 5;
vector<int> G[MAX_N];
segment_tree* dp[MAX_N];
int max_depth[MAX_N];
int N, D;

int DFS(int u, int d = 0) {
  max_depth[u] = d;
  dp[u] = new segment_tree(0, 2 << __lg(N));
  int ans = 1;
  for (int v : G[u]) {
    DFS(v, d + 1);
    ans += dp[v]->query(d + D);
    if (max_depth[u] < max_depth[v]) {
      swap(max_depth[u], max_depth[v]);
      swap(dp[u], dp[v]);
    }
    vector<int> vals;
    for (int i = d + 1; i <= max_depth[v]; i++) {
      vals.push_back(dp[u]->query(max(i, 2 * d + D - i)));
    }
    for (int i = d + 1; i <= max_depth[v]; i++) {
      dp[u]->modify(i, dp[u]->get(i) + dp[v]->query(max(i, 2 * d + D - i)));
    }
    for (int i = d + 1, j = 0; i <= max_depth[v]; i++) {
      dp[u]->modify(i, dp[v]->get(i) + vals[j++]);
    }
  }
  dp[u]->modify(d, ans);
  return dp[u]->query(0);
}

int main() {
  cin.tie(nullptr)->sync_with_stdio(false);
  cin >> N >> D;
  if (D == 1) {
    cout << N << "\n";
    exit(0);
  }
  for (int i = 1; i < N; i++) {
    int p;
    cin >> p;
    G[p].push_back(i);
  }
  cout << DFS(0) << "\n";
}

• Subtask 1: 11 / 11
• Subtask 2: 40 / 40
• Subtask 3: 49 / 49