제출 #980947

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
980947		RedGrey1993	JOI tour (JOI24_joitour)	C++17	100 / 100	607 ms	204336 KiB

joitour

#include "joitour.h"
#include <bits/stdc++.h>

using namespace std;
 
template <typename T1, typename T2> istream &operator>>(istream &is, pair<T1, T2> &pa) { is >> pa.first >> pa.second; return is; }
template <typename T> istream &operator>>(istream &is, vector<T> &vec) { for (auto &v : vec) is >> v; return is; }
template <typename T1, typename T2> ostream &operator<<(ostream &os, const pair<T1, T2> &pa) { os << "(" << pa.first << "," << pa.second << ")"; return os; }
template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) { os << "["; for (auto v : vec) os << v << ","; os << "]"; return os; }
template <typename T> ostream &operator<<(ostream &os, const deque<T> &vec) { os << "deq["; for (auto v : vec) os << v << ","; os << "]"; return os; }
template <typename T> ostream &operator<<(ostream &os, const set<T> &vec) { os << "{"; for (auto v : vec) os << v << ","; os << "}"; return os; }
template <typename T> ostream &operator<<(ostream &os, const multiset<T> &vec) { os << "{"; for (auto v : vec) os << v << ","; os << "}"; return os; }
template <typename T> ostream &operator<<(ostream &os, const unordered_set<T> &vec) { os << "{"; for (auto v : vec) os << v << ","; os << "}"; return os; }
template <typename T> ostream &operator<<(ostream &os, const unordered_multiset<T> &vec) { os << "{"; for (auto v : vec) os << v << ","; os << "}"; return os; }
template <typename TK, typename TV> ostream &operator<<(ostream &os, const unordered_map<TK, TV> &mp) { os << "{"; for (auto v : mp) os << v.first << "=>" << v.second << ","; os << "}"; return os; }
template <typename TK, typename TV> ostream &operator<<(ostream &os, const map<TK, TV> &mp) { os << "{"; for (auto v : mp) os << v.first << "=>" << v.second << ","; os << "}"; return os; }
template <typename T> void resize_array(vector<T> &vec, int len) { vec.resize(len); }
template <typename T, typename... Args> void resize_array(vector<T> &vec, int len, Args... args) { vec.resize(len); for (auto &v : vec) resize_array(v, args...); }
template <typename T1, typename T2> pair<T1, T2> operator+(const pair<T1, T2> &l, const pair<T1, T2> &r) { return make_pair(l.first + r.first, l.second + r.second); }
template <typename T1, typename T2> pair<T1, T2> operator-(const pair<T1, T2> &l, const pair<T1, T2> &r) { return make_pair(l.first - r.first, l.second - r.second); }
long long gcd(long long a, long long b) { return b ? gcd(b, a % b) : a; }
mt19937 mrand(random_device{}());
int rnd(int x) { return mrand() % x; }
 
#define rep(i, a, n) for (int i = a; i < (n); i++)
#define per(i, a, n) for (int i = (n)-1; i >= a; i--)
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define fi first
#define se second
#define sz(x) ((int)(x).size())
typedef vector<int> vi;
typedef long long ll;
typedef unsigned int uint;
typedef unsigned long long ull;
typedef pair<int, int> pii;
typedef double db;
#if DEBUG
#define dbg(x) cerr << #x << " = " << (x) << " (L" << __LINE__ << ") " << __FILE__ << endl;
#else
#define dbg(x)
#endif

template <typename T>
class StaticTopTree {
 public:
  enum Type { Vertex, Compress, Rake, AddEdge, AddVertex };
  StaticTopTree(vector<vector<int>> &g, int root = 0) : g_(g), root_(root) {
    int n = g_.size();
    parent_.resize(4 * n, -1);
    left_.resize(4 * n, -1);
    right_.resize(4 * n, -1);
    vertex_type_.resize(4 * n, Type::Vertex);
    paths_.resize(4 * n);
    points_.resize(4 * n);
    cur_idx_ = n;
    build();
  }

 private:
  // do HLD(Heavy Light Decomposition), move heavy edge to g_[c][0];
  int dfs(int c) {
    int s = 1, best = 0;
    for (int &d : g_[c]) {
      int t = dfs(d);
      s += t;
      if (best < t) best = t, swap(d, g_[c][0]);
    }
    return s;
  }
  int add(int k, int l, int r, Type t) {
    if (k == -1) k = cur_idx_++;
    parent_[k] = -1, left_[k] = l, right_[k] = r, vertex_type_[k] = t;
    if (l != -1) parent_[l] = k;
    if (r != -1) parent_[r] = k;
    return k;
  }
  pair<int, int> merge(const vector<pair<int, int>> &a, Type t) {
    if (a.size() == 1) return a[0];
    int u = 0;
    for (auto &[_, s] : a) u += s;
    vector<pair<int, int>> b, c;
    // make sure b/c at least has 1 item
    for (auto& [i, s] : a) (u > s ? b : c).emplace_back(i, s), u -= s * 2;
    // if use (u > 0 ? b : c), c may be 0 size, will be RE
    // for (auto &[i, s] : a) (u > 0 ? b : c).emplace_back(i, s), u -= s * 2;
    auto [i, si] = merge(b, t);
    auto [j, sj] = merge(c, t);
    return {add(-1, i, j, t), si + sj};
  }
  pair<int, int> compress(int i) {
    vector<pair<int, int>> chs{add_vertex(i)};
    while (!g_[i].empty()) chs.push_back(add_vertex(i = g_[i][0]));
    return merge(chs, Type::Compress);
  }
  pair<int, int> rake(int i) {
    vector<pair<int, int>> chs;
    for (int j = 1; j < (int)g_[i].size(); j++)
      chs.push_back(add_edge(g_[i][j]));
    return chs.empty() ? make_pair(-1, 0) : merge(chs, Type::Rake);
  }
  pair<int, int> add_edge(int i) {
    auto [j, sj] = compress(i);
    return {add(-1, j, -1, Type::AddEdge), sj};
  }
  pair<int, int> add_vertex(int i) {
    auto [j, sj] = rake(i);
    return {add(i, j, -1, j == -1 ? Type::Vertex : Type::AddVertex), sj + 1};
  }
  void build() {
    dfs(root_);
    auto [i, n] = compress(root_);
    stt_root_ = i;
  }
  // private:
 public:
  // adjacent list of initial tree, chridren are stored in g_[parent],
  // parent can not be stored in g_[child] in this implement!!!
  // g_ must be a rooted tree
  vector<vector<int>> &g_;
  int root_;      // an index of the root in g
  int stt_root_;  // an index of the root in static top tree
  vector<int> parent_, left_, right_;  // parent, left child, right child
  vector<Type> vertex_type_;           // type of vertices
  int cur_idx_;                        // a variable for the index increment

  vector<int> values;

 private:
  // u: upper boundary vertex, v: lower boundary vertex
  struct TP {
    // c01u -> u can be 1, cu1v -> u or v can be 1
    // c0_u_2 -> '_' means u can not be 0/2
    ll c012,c01u,c2,c0,cu12,c01v,cv12,cu1v,c0_u_2,cvu_2,c0_uv;
    friend ostream &operator<<(ostream &os, const TP &p) { 
      return os << "c012:" << p.c012 << "\n"
                << "c01u:" << p.c01u << "\n"
                << "c2:" << p.c2 << "\n"
                << "c0:" << p.c0 << "\n"
                << "cu12:" << p.cu12 << "\n"
                << "c01v:" << p.c01v << "\n"
                << "cv12:" << p.cv12 << "\n"
                << "cu1v:" << p.cu1v << "\n"
                << "c0_u_2:" << p.c0_u_2 << "\n"
                << "cvu_2:" << p.cvu_2 << "\n"
                << "c0_uv:" << p.c0_uv << "\n"; 
    }
  };
  using Point = TP;
  using Path = TP;
  vector<Point> points_;
  vector<Path> paths_;
  Path vertex(int i) { 
    TP tp;
    tp.c012 = 0;
    tp.c01u = 0;
    tp.c2 = (values[i] == 2);
    tp.c0 = (values[i] == 0);
    tp.cu12 = 0;
    tp.c01v = 0;
    tp.cv12 = 0;
    tp.cu1v = (values[i] == 1);
    tp.c0_u_2 = 0;
    tp.cvu_2 = 0;
    tp.c0_uv = 0;
    return tp; 
  }
  Path compress(Path p, Path c) { 
    TP tp;
    tp.c012 = p.c012 + c.c012 + p.c01v * c.c2 + p.cv12 * c.c0 + p.c0 * c.cu12 + p.c2 * c.c01u;
    tp.c01u = p.c01u + c.c01u + p.cu1v * c.c0;
    tp.c2 = p.c2 + c.c2;
    tp.c0 = p.c0 + c.c0;
    tp.cu12 = p.cu12 + c.cu12 + p.cu1v * c.c2;
    tp.c01v = p.c01v + c.c01v + p.c0 * c.cu1v;
    tp.cv12 = p.cv12 + c.cv12 + p.c2 * c.cu1v;
    tp.cu1v = p.cu1v + c.cu1v;
    tp.c0_u_2 = p.c0_u_2 + c.c0 * p.cvu_2 + c.c2 * p.c0_uv;
    tp.cvu_2 = p.cvu_2;
    tp.c0_uv = p.c0_uv;
    return tp;
  }
  Path add_vertex(Point p, int i) { 
    TP tp;
    // tp.fv = values[i];
    tp.c012 = p.c012 + p.cu12 * (values[i] == 0) + p.c01u * (values[i] == 2) + p.c0_u_2 * (values[i] == 1);
    tp.c01u = p.c01u + p.c0 * (values[i] == 1);
    tp.c2 = p.c2 + (values[i] == 2);
    tp.c0 = p.c0 + (values[i] == 0);
    tp.cu12 = p.cu12 + (values[i] == 1) * p.c2;
    tp.c01v = tp.c01u;
    tp.cv12 = tp.cu12;
    tp.cu1v = (values[i] == 1);
    tp.c0_u_2 = p.c0_u_2;
    tp.cvu_2 = p.c2;
    tp.c0_uv = p.c0;
    return tp;
  }
  Point add_edge(Path p) {
    TP tp = p;
    tp.cu1v = 0;
    tp.c0_u_2 = 0;
    tp.cvu_2 = 0;
    tp.c0_uv = 0;
    return tp;
  }
  Point rake(Point p, Point c) {
    TP tp;
    // tp.fv = values[i];
    tp.c012 = p.c012 + c.c012 + p.c01u * c.c2 + p.cu12 * c.c0 + p.c2 * c.c01u + p.c0 * c.cu12;
    tp.c01u = p.c01u + c.c01u;
    tp.c2 = p.c2 + c.c2;
    tp.c0 = p.c0 + c.c0;
    tp.cu12 = p.cu12 + c.cu12;
    tp.c01v = 0;
    tp.cv12 = 0;
    tp.cu1v = 0;
    tp.c0_u_2 = p.c0_u_2 + c.c0_u_2 + p.c0 * c.c2 + p.c2 * c.c0;
    // tp.cvu_2 = 0;
    // tp.c0_uv = 0;
    return tp;
  }

  void update_node(int k) {
    if (vertex_type_[k] == Vertex) paths_[k] = vertex(k);
    else if (vertex_type_[k] == Compress) paths_[k] = compress(paths_[left_[k]], paths_[right_[k]]);
    else if (vertex_type_[k] == AddVertex) paths_[k] = add_vertex(points_[left_[k]], k);
    else if (vertex_type_[k] == AddEdge) points_[k] = add_edge(paths_[left_[k]]);
    else if (vertex_type_[k] == Rake) points_[k] = rake(points_[left_[k]], points_[right_[k]]);
  }

  void init(int k) {
    if (left_[k] != -1) init(left_[k]);
    if (right_[k] != -1) init(right_[k]);
    update_node(k);
    dbg(mp(k, vertex_type_[k]));
    dbg(points_[k]);
    dbg(paths_[k]);
  }

 public:
  void update_tree(int k, T v) {
    values[k] = v;
    while (k != -1) {
      update_node(k);
      k = parent_[k];
    }
  }

  void init() {
    init(stt_root_);
  }

  T get_answer() {
    return paths_[stt_root_].c012;
  }
};

unique_ptr<StaticTopTree<ll>> p_stt = nullptr;

void init(int N, std::vector<int> F, std::vector<int> U, std::vector<int> V, int Q) {
  vector<vi> edges(N);

  rep(i,0,N-1) {
    edges[U[i]].emplace_back(V[i]);
    edges[V[i]].emplace_back(U[i]);
  }
  vector<vi> g(N);
  function<void(int,int)> dfs=[&] (int u, int p) {
    for(int v : edges[u]) {
      if (v != p) {
        g[u].emplace_back(v);
        dfs(v, u);
      }
    }
  };
  dfs(0, -1);
  dbg(g);

  p_stt = make_unique<StaticTopTree<ll>>(g);
  p_stt->values = F;
  dbg(p_stt->values);
  p_stt->init();
}

void change(int X, int Y) {
  p_stt->update_tree(X, Y);
}

long long num_tours() {
  return p_stt->get_answer();
}

• Subtask 1: 6 / 6
• Subtask 2: 8 / 8
• Subtask 3: 6 / 6
• Subtask 4: 16 / 16
• Subtask 5: 16 / 16
• Subtask 6: 34 / 34
• Subtask 7: 14 / 14