Submission #229692

#	Time	Username	Problem	Language	Result	Execution time	Memory
229692		xiaowuc1	트리 (KOI16_tree)	C++14	100 / 100	481 ms	51452 KiB

tree

#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <complex>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <vector>

using namespace std;

// BEGIN NO SAD
#define rep(i, a, b) for(int i = a; i < (b); ++i)
#define trav(a, x) for(auto& a : x)
#define all(x) x.begin(), x.end()
#define sz(x) (int)(x).size()
#define derr if(1) cerr
typedef vector<int> vi;
// END NO SAD

typedef long long ll;
typedef pair<int, int> pii;
typedef vector<vector<ll>> matrix;
typedef pair<int, pii> state;

const int KOOSAGA_DEPTH = 18;
const int KOOSAGA_SZ = 1 << KOOSAGA_DEPTH;
bool koosagatree[2 * KOOSAGA_SZ];
void upd(int idx) {
  idx += KOOSAGA_SZ;
  while(idx) {
    koosagatree[idx] = true;
    idx /= 2;
  }
}
bool qry(int lhs, int rhs) {
  lhs += KOOSAGA_SZ;
  rhs += KOOSAGA_SZ;
  while(lhs <= rhs) {
    if(lhs == rhs) return koosagatree[lhs];
    if(lhs%2) if(koosagatree[lhs++]) return true;
    if(rhs%2==0) if(koosagatree[rhs--]) return true;
    lhs /= 2;
    rhs /= 2;
  }
  return false;
}

vector<int> children[KOOSAGA_SZ];
int vertextoedge[KOOSAGA_SZ]; // vertextoedge[i] gives koosaga tree ID for (parent of i, i) edge
int treedepth[KOOSAGA_SZ];
int treetop[KOOSAGA_SZ];
int treesz[KOOSAGA_SZ];
int lca[KOOSAGA_SZ][KOOSAGA_DEPTH];

bool _sort_by_sz(int a, int b) {
  return treesz[a] > treesz[b];
}
void dfsforhld(int curr, int top, int& id) {
  treetop[curr] = top;
  if(sz(children[curr]) == 0) return;
  sort(all(children[curr]), _sort_by_sz);
  vertextoedge[children[curr][0]] = id++;
  dfsforhld(children[curr][0], top, id);
  for(int i = 1; i < sz(children[curr]); i++) {
    vertextoedge[children[curr][i]] = id++;
    dfsforhld(children[curr][i], children[curr][i], id);
  }
}
void dfsforsz(int curr) {
  treesz[curr] = 1;
  for(int out: children[curr]) {
    treedepth[out] = treedepth[curr] + 1;
    dfsforsz(out);
    treesz[curr] += treesz[out];
  }
}

int getlca(int a, int b) {
  if(treedepth[a] < treedepth[b]) swap(a, b);
  for(int d = KOOSAGA_DEPTH-1; d >= 0; d--) {
    if(treedepth[a] - (1<<d) >= treedepth[b]) a = lca[a][d];
  }
  assert(treedepth[a] == treedepth[b]);
  for(int d = KOOSAGA_DEPTH-1; d >= 0; d--) {
    if(lca[a][d] != lca[b][d]) {
      a = lca[a][d];
      b = lca[b][d];
    }
  }
  while(a != b) {
    a = lca[a][0];
    b = lca[b][0];
  }
  return a;
}

bool pathqry(int root, int curr) {
  assert(treedepth[root] <= treedepth[curr]);
  while(treedepth[curr] > treedepth[root]) {
    assert(curr);
    if(curr == treetop[curr]) {
      if(qry(vertextoedge[curr], vertextoedge[curr])) return false;
      curr = lca[curr][0];
      continue;
    }
    if(treedepth[treetop[curr]] >= treedepth[root]) {
      // go all the way to top[curr]
      int nume = treedepth[curr] - treedepth[treetop[curr]] - 1;
      assert(nume >= 0);
      if(qry(vertextoedge[curr]-nume, vertextoedge[curr])) return false;
      curr = treetop[curr];
      continue;
    }
    // go all the way to root
    int nume = treedepth[curr] - treedepth[root] - 1;
    assert(nume >= 0);
    if(qry(vertextoedge[curr]-nume, vertextoedge[curr])) return false;
    return true;
  }
  return true;
}

bool connected(int a, int b) {
  int r = getlca(a, b);
  assert(r);
  return pathqry(r, a) && pathqry(r, b);
}

void solve() {
  int n, q;
  cin >> n >> q;
  for(int i = 2; i <= n; i++) {
    cin >> lca[i][0];
    children[lca[i][0]].push_back(i);
  }
  dfsforsz(1);
  for(int d = 1; d < KOOSAGA_DEPTH; d++) {
    for(int i = 1; i <= n; i++) {
      lca[i][d] = lca[lca[i][d-1]][d-1];
    }
  }
  {
    int id = 0;
    dfsforhld(1, 1, id);
  }
  while(q--) {
    int a, b, t;
    cin >> a >> b >> t;
    assert(t == 0 || t == 1);
    if(connected(a, b)) {
      cout << "YES\n";
      if(t==1 && a > 1) upd(vertextoedge[a]);
    }
    else {
      cout << "NO\n";
      if(t==1 && b > 1) upd(vertextoedge[b]);
    }
  }
}
// are there edge cases (N=1?)
// are array sizes proper (scaled by proper constant, for example 2* for koosaga tree)
// integer overflow?
// DS reset properly between test cases

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL); cout.tie(NULL);
  solve();
}

• Subtask 1: 9 / 9
• Subtask 2: 13 / 13
• Subtask 3: 14 / 14
• Subtask 4: 19 / 19
• Subtask 5: 45 / 45