Submission #774762

#	Time	Username	Problem	Language	Result	Execution time	Memory
774762		arush_agu	Ball Machine (BOI13_ballmachine)	C++17	100 / 100	597 ms	27204 KiB

ballmachine

#include <algorithm>
#include <cctype>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <functional>
#include <iomanip>
#include <iostream>
#include <iterator>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#ifdef DEBUG
#include <time.h>
#endif

#define all(a) (a).begin(), (a).end()
#define rev(a) (a).rbegin(), (a).rend()
#define F first
#define S second
int recur_depth = 0;
#ifdef DEBUG
#define dbg(x)                                                                 \
  {                                                                            \
    ++recur_depth;                                                             \
    auto x_ = x;                                                               \
    --recur_depth;                                                             \
    cerr << string(recur_depth, '\t') << "\e[91m" << __func__ << ":"           \
         << __LINE__ << "\t" << #x << " = " << x_ << "\e[39m" << endl;         \
  }
#else
#define dbg(x)
#endif

using namespace std;
using namespace __gnu_pbds;

typedef pair<int, int> ii;

typedef long long ll;
typedef long double ld;
typedef pair<ll, ll> llll;

typedef vector<int> vi;
typedef vector<vi> vvi;
typedef vector<pair<int, int>> vii;
typedef vector<pair<int, ii>> viii;
typedef vector<vii> vvii;
typedef vector<viii> vviii;

typedef vector<ll> vl;
typedef vector<vl> vvl;
typedef vector<pair<ll, ll>> vll;
typedef vector<vll> vvll;

typedef vector<bool> vb;

template <class type1>
using ordered_set = tree<type1, null_type, less<type1>, rb_tree_tag,
                         tree_order_statistics_node_update>;

template <typename A, typename B>
ostream &operator<<(ostream &os, const pair<A, B> &p) {
  return os << '(' << p.first << ", " << p.second << ')';
}
template <typename T_container, typename T = typename enable_if<
                                    !is_same<T_container, string>::value,
                                    typename T_container::value_type>::type>
ostream &operator<<(ostream &os, const T_container &v) {
  os << '{';
  string sep;
  for (const T &x : v)
    os << sep << x, sep = ", ";
  return os << '}';
}

const ll MOD = 1e9 + 7;
// const ll MOD = 998244353;
const ll INF = 1e9;
const ld EPS = 1e-9;

const int MAXLG = 18;
const int MAXN = 1e5 + 1;
int n, q, root;
int par[MAXN], mn_subtree[MAXN], depth[MAXN], tin[MAXN], tout[MAXN],
    up[MAXLG][MAXN], cnt_cball[MAXN];
vi child[MAXN];
bool is_ball[MAXN], inq[MAXN];
int timer = 0;

// void f(int u, int p) {
//   tin[u] = ++timer;
//   mn_subtree[u] = u;
//   for (int &v : child[u]) {
//     depth[v] = depth[u] + 1;
//     f(v, u);
//     mn_subtree[u] = min(mn_subtree[u], mn_subtree[v]);
//   }
//   tout[u] = ++timer;
// }

void solve() {
  cin >> n >> q;
  for (int i = 0; i < n; ++i) {
    cin >> par[i];
    if (--par[i] == -1)
      root = i;
  }

  for (int i = 0; i < n; ++i)
    if (i != root)
      child[par[i]].push_back(i);

  depth[root] = 0;
  function<void(int, int)> f = [&](int u, int p) {
    tin[u] = ++timer;
    mn_subtree[u] = u;
    for (int &v : child[u]) {
      depth[v] = depth[u] + 1;
      f(v, u);
      mn_subtree[u] = min(mn_subtree[u], mn_subtree[v]);
    }
    tout[u] = ++timer;
  };
  f(root, -1);

  for (int i = 0; i < n; ++i)
    up[0][i] = par[i];

  for (int j = 1; j < MAXLG; ++j)
    for (int i = 0; i < n; ++i) {
      if (up[j - 1][i] != -1)
        up[j][i] = up[j - 1][up[j - 1][i]];
      else
        up[j][i] = -1;
    }

  auto is_anc = [&](int u, int v) {
    return tin[v] <= tin[u] && tout[u] <= tout[v];
  };

  auto kth_anc = [&](int u, int k) {
    for (int i = MAXLG - 1; i >= 0; --i)
      if (k & (1 << i)) {
        u = up[i][u];
      }
    return u;
  };

  auto just_lca = [&](int u, int v) -> ii {
    if (depth[u] > depth[v])
      u = kth_anc(u, depth[u] - depth[v]);
    else
      v = kth_anc(v, depth[v] - depth[u]);

    u = kth_anc(u, depth[u] - depth[v]);

    for (int i = MAXLG - 1; i >= 0; --i)
      if (up[i][u] != up[i][v]) {
        u = up[i][u];
        v = up[i][v];
      }

    return {u, v};
  };

  auto s_cmp = [&](int u, int v) -> bool {
    if (u == v)
      return 0;
    if (is_anc(u, v))
      return 0;
    if (is_anc(v, u))
      return 1;
    auto [x, y] = just_lca(u, v);
    return mn_subtree[x] > mn_subtree[y];
  };

  priority_queue<int, vi, decltype(s_cmp)> s(s_cmp);
  for (int i = 0; i < n; ++i)
    if (child[i].empty())
      s.push(i), inq[i] = 1;

  while (q--) {
    int t, x;
    cin >> t >> x;

    if (t == 1) {
      for (int i = 1; i <= x; ++i) {
        int u = s.top();
        s.pop();

        inq[u] = 0;
        is_ball[u] = 1;

        if (u != root)
          if (++cnt_cball[par[u]] == child[par[u]].size() && !inq[par[u]])
            s.push(par[u]), inq[par[u]] = 1;

        if (i == x)
          cout << u + 1 << "\n";
      }
    } else {
      --x;
      int ans = x;
      for (int i = MAXLG - 1; i >= 0; i--)
        if (up[i][ans] != -1 && is_ball[up[i][ans]])
          ans = up[i][ans];

      is_ball[ans] = 0;
      if (ans != root)
        cnt_cball[par[ans]]--;
      if (!inq[ans])
        s.push(ans), inq[ans] = 1;

      cout << depth[x] - depth[ans] << "\n";
    }
  }
}

int main() {
  ios_base::sync_with_stdio(0);
  cin.tie(NULL);

  clock_t start = clock();

  int test_cases = 1;
  // cin >> test_cases;

  while (test_cases--)
    solve();

#ifdef DEBUG
  cerr << fixed << setprecision(10)
       << "\nTime Taken: " << (double)(clock() - start) / CLOCKS_PER_SEC
       << "s\n";
#endif
  return 0;
}

Compilation message (stderr)

ballmachine.cpp: In function 'void solve()':
ballmachine.cpp:207:35: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  207 |           if (++cnt_cball[par[u]] == child[par[u]].size() && !inq[par[u]])
      |               ~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~
ballmachine.cpp: In function 'int main()':
ballmachine.cpp:235:11: warning: unused variable 'start' [-Wunused-variable]
  235 |   clock_t start = clock();
      |           ^~~~~

• Subtask 1: 25 / 25
• Subtask 2: 30 / 30
• Subtask 3: 40 / 40
• Subtask 4: 5 / 5