Submission #774107

#	Time	Username	Problem	Language	Result	Execution time	Memory
774107		arush_agu	Ball Machine (BOI13_ballmachine)	C++17	46.72 / 100	1090 ms	27140 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 + 10;
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];

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;
  int timer = 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) {
      up[j][i] = -1;
      if (up[j - 1][i] != -1)
        up[j][i] = up[j - 1][up[j - 1][i]];
    }

  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 {
    bool swaped = 0;
    if (depth[u] < depth[v])
      swaped = 1, swap(u, v);

    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];
      }

    if (swaped)
      return {v, u};
    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];
    // return mn_subtree[x] < mn_subtree[y];
  };

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

  auto print_pq = [&]() {
    vi tmp;
    while (s.size()) {
      tmp.push_back(s.top());
      s.pop();
    }

    cerr << tmp << "\n";

    for (int x : tmp)
      s.push(x);
  };

  // print_pq();

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

    // cerr << t << " " << x << "\n";
    // // cerr << "\tIS BAll: " << is_ball << "\n";
    // cerr << "\tPQ: ";
    // print_pq();

    if (t == 1) {
      for (int i = 1; i <= x; ++i) {
        // cerr << "\t\tI: " << i << "\n";
        // cerr << "\t\t\tPQ: ";
        // print_pq();
        // cerr << "\t\t\tInq: {";
        // for (int i = 0; i < n; i++)
        //   cerr << inq[i] << ",";
        // cerr << "}\n";

        // int u = s.top();
        // s.pop();
        // auto it = s.begin();
        // int u = *it;
        // s.erase(it);
        int u = s.top();
        s.pop();

        inq[u] = 0;

        // cerr << "\tU: " << u << "\n";

        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;
          // s.insert(par[u]);
        }

        if (i == x)
          cout << u + 1 << "\n";
      }

      // cerr << "\tIB:";
      // for (int i = 0; i < n; i++)
      //   cerr << is_ball[i] << ",";
      // cerr << "\n";
    } else {
      --x;
      // cerr << "\t" << x << "\n";
      // is_ball[x] = 0;
      // --cnt_cball[par[x]];
      // int res = 0;
      // while (x != root && is_ball[par[x]]) {
      //   is_ball[x] = 1;
      //   ++cnt_cball[par[x]];
      //   x = par[x];
      //   is_ball[x] = 0;
      //   if (x != root)
      //     --cnt_cball[par[x]];
      //   res++;
      // }

      // cerr << "\t\t" << res << "\t" << x << "\n";
      // cerr << "\t\t" << s.count(x) << "\n";
      // s.insert(x);
      // cerr << "\t\t";
      // print_pq();

      // if (!inq[x])
      //   s.push(x), inq[x] = 1;
      // if (x != root) {
      //   auto it = s.find(par[x]);
      //   if (it != s.end())
      //     s.erase(it);
      // }

      int l = 1, r = depth[x], ans = 0;
      while (l <= r) {
        int mid = (l + r) / 2;
        if (is_ball[kth_anc(x, mid)]) {
          ans = mid;
          l = mid + 1;
        } else
          r = mid - 1;
      }

      ans = kth_anc(x, ans);

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

      cout << depth[x] - depth[ans] << "\n";
    }
    // cerr << "\t" << is_ball << "\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:234:35: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  234 |           if (++cnt_cball[par[u]] == child[par[u]].size() && !inq[par[u]])
      |               ~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~
ballmachine.cpp:186:8: warning: variable 'print_pq' set but not used [-Wunused-but-set-variable]
  186 |   auto print_pq = [&]() {
      |        ^~~~~~~~
ballmachine.cpp: In function 'int main()':
ballmachine.cpp:305:11: warning: unused variable 'start' [-Wunused-variable]
  305 |   clock_t start = clock();
      |           ^~~~~

• Subtask 1: 5.769230769230769 / 25
• Subtask 2: 10.714285714285714 / 30
• Subtask 3: 28.571428571428573 / 40
• Subtask 4: 1.6666666666666667 / 5