Submission #46498682

---

Source Code Expand

#include "conveyor.h"
#include <bits/stdc++.h>
#define pii pair<int, int>
#define mp make_pair 

using std::vector;
using std::pair;

const int N = 1e5 + 10;
int n, fa[N], to[N], dep[N];
int cnt[N], id[N];
vector<pii> G[N];
vector<int> T[3];

void dfs(int u) {
  T[dep[u] % 3].emplace_back(u);
  for(auto x : G[u]) {
    int v = x.first;
    if(v == fa[u]) continue;
    dep[v] = dep[u] + 1;
    id[v] = x.second, fa[v] = u;
    dfs(v);
  }
}

int count(int x) {
  int res = 0;
  for(int i : T[x])
    res += (cnt[i] > 0);
  return res;
}

vector<int> query(vector<int> vec) {
  static int vis[N];
  for(int i = 0; i < n; ++i)
    vis[i] = 0;
  for(int x : vec) vis[x] = 1;
  vector<int> E(n - 1, 0);
  for(int i = 0; i < n - 1; ++i)
    if(~to[i] && vis[to[i]]) E[i] = 1;
  vector<int> pos(n, 0);
  for(int i : vec) pos[i] = 1;
  return Query(E, vec);
}

void Solve(int _n, vector<int> A, vector<int> B) {
  n = _n;
  for(int i = 0; i < n - 1; ++i) {
    G[A[i]].emplace_back(pii{B[i], i});
    G[B[i]].emplace_back(pii{A[i], i});
    to[i] = -1, ++cnt[A[i]], ++cnt[B[i]];
  }

  auto set = [&](int x, int y) {
    --cnt[A[x]], --cnt[B[x]];
    to[x] = y;
  };

  dfs(1);
  while(true) {
    int t = 0;
    for(int i = 1; i <= 2; ++i)
      if(count(i) > count(t)) t = i;

    vector<int> pos;
    for(int i : T[t]) 
      if(cnt[i] > 0) pos.emplace_back(i);

    if(!pos.size()) {
      vector<int> ans(n - 1, 0);
      for(int i = 0; i < n - 1; ++i)
        if(B[i] != to[i]) ans[i] = 1;
      Answer(ans);
      return ;
    }

    vector<int> res = query(pos);
    static int vis[N];
    for(int x : pos) vis[x] = 0;

    for(int x : T[(t + 1) % 3]) 
      if(res[x]) set(id[x], x), vis[fa[x]] = 1;

    for(int x : pos) {
      if(res[x]) {
        for(auto i : G[x])
          if(to[i.second] != -1) set(i.second, x);
      } else if(!vis[x]) {
        set(id[x], fa[x]);
      }
    }
  }
}

Submission Info

Judge Result