Submission #774010

#	Time	Username	Problem	Language	Result	Execution time	Memory
774010		arush_agu	Ball Machine (BOI13_ballmachine)	C++17	42.11 / 100	1089 ms	26124 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];
vi child[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 j = 0; j < MAXLG; j++)
    for (int i = 0; i < n; i++)
      up[j][i] = -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]];

  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 1;
    if (is_anc(v, u))
      return 0;
    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.insert(i);

  auto print_pq = [&]() {
    cerr << s << "\n";
    // vi tmp;
    // while (!s.empty()) {
    //   tmp.push_back(s.top());
    //   s.pop();
    // }
    // cerr << tmp << "\n";
    // for (int x : tmp)
    //   s.push(x);
  };

  // cerr << "PQ: ";
  // print_pq();

  vb is_ball(n);
  vi cnt_cball(n);
  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++) {
        // int u = s.top();
        // s.pop();
        int u = *s.begin();
        s.erase(s.begin());

        // cerr << "\tU: " << u << "\n";
        assert(is_ball[u] == 0);

        is_ball[u] = 1;
        if (u != root) {
          if (++cnt_cball[par[u]] == child[par[u]].size())
            // s.push(par[u]);
            s.insert(par[u]);
        }

        if (i == x)
          cout << u + 1 << "\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();

      // s.push(x);
      if (x != root) {
        if (s.count(par[x]))
          s.erase(s.find(par[x]));
      }
      cout << res << "\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:223:35: warning: comparison of integer expressions of different signedness: '__gnu_cxx::__alloc_traits<std::allocator<int>, int>::value_type' {aka 'int'} and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  223 |           if (++cnt_cball[par[u]] == child[par[u]].size())
      |               ~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~
ballmachine.cpp:185:8: warning: variable 'print_pq' set but not used [-Wunused-but-set-variable]
  185 |   auto print_pq = [&]() {
      |        ^~~~~~~~
ballmachine.cpp: In function 'int main()':
ballmachine.cpp:268:11: warning: unused variable 'start' [-Wunused-variable]
  268 |   clock_t start = clock();
      |           ^~~~~

• Subtask 1: 21.153846153846153 / 25
• Subtask 2: 19.285714285714285 / 30
• Subtask 3: 0 / 40
• Subtask 4: 1.6666666666666667 / 5