dsu/BinaryLifting.cpp

• View this file on GitHub
• Last update: 2023-02-27 10:03:35-03:00

Depends on

• bits/stdc++.h
• dsu/SpanningTree.cpp
• utils/LazyArray.cpp

Code

#ifndef _LIB_DSU_BINARY_LIFTING
#define _LIB_DSU_BINARY_LIFTING
#include <bits/stdc++.h>
#include "SpanningTree.cpp"

namespace lib {
using namespace std;
namespace dsu {

template<typename D>
struct BinaryLifting : public D {
  using D::parent;
  vector<vector<int>> P;
  int K;

  BinaryLifting() : D() {}
  BinaryLifting(int n) : D(n) {
    P = decltype(P)(n, vector<int>(__lg(n)+1, -1));
    K = __lg(n)+1;
  }
  virtual void clear() override {
    D::clear();
    int n = P.size();
    P = decltype(P)(n, vector<int>(K, -1));
  }
  virtual int merge(int u, int v) override {
    if(!D::merge(u, v)) return 0;
    this->traverse_last_small([this](int u, int p, vector<int>&) {
      for(int& x : P[u]) x = -1;
      P[u][0] = p;
      for(int i = 1; i < K; i++) {
        if(P[u][i-1] == -1) break;
        P[u][i] = P[P[u][i-1]][i-1];
      }
    }, no_op_visitor);
    return 1;
  }
  int parent(int u, int k) {
    assert(k >= 0);
    for(int i = K-1; i >= 0; i--) {
      if(!((k>>i)&1)) continue;
      u = P[u][i];
      if(u == -1) return -1;
    }
    return u;
  }
};
} // namespace dsu
} // namespace lib

#endif

#line 1 "dsu/BinaryLifting.cpp"


#include <bits/stdc++.h>
#line 1 "dsu/SpanningTree.cpp"


#line 1 "utils/LazyArray.cpp"


#line 4 "utils/LazyArray.cpp"

namespace lib {
  using namespace std;
template<typename T>
struct LazyArray {
  vector<T> v;
  vector<int> lz;
  T lz_v;
  int tempo = 1;

  LazyArray() {}
  LazyArray(int n, T lz_v) : v(n), lz(n), lz_v(lz_v) {}

  int size() const { return v.size(); }
  void fill(T vv) { lz_v = vv, tempo++; }
  void clear() {
    fill(T());
  }
  T get(int i) const {
    return tempo == lz[i] ? v[i] : lz_v;
  }
  T operator[](int i) const {
    return get(i);
  }
  T& operator[](int i) {
    if(lz[i] != tempo) {
      lz[i] = tempo;
      v[i] = lz_v;
    }
    return v[i];
  }
};
} // namespace lib


#line 5 "dsu/SpanningTree.cpp"

namespace lib {
using namespace std;
namespace dsu {

const auto no_op_visitor = [](int, int, const vector<int>&) -> void {};

template<typename D>
struct SpanningTree : public D {
  using D::last_swapped_;

  vector<vector<int>> adj;
  vector<int> pai, depth;
  LazyArray<char> vis;
  pair<int, int> last_edge_;

  SpanningTree() : D() {}
  SpanningTree(int n) : D(n), adj(n), pai(n, -1), vis(n, 0), depth(n, 0) {}
  virtual void clear() override {
    D::clear();
    for(int i = 0; i < adj.size(); i++)
      adj[i].clear();
    fill(pai.begin(), pai.end(), -1);
    fill(depth.begin(), depth.end(), 0);
    vis.clear();
    last_edge_ = {-1, -1};
  }
  virtual int merge(int u, int v) override {
    if(!D::merge(u, v)) return 0;
    if(last_swapped_)
      swap(u, v);
    last_edge_ = {u, v};
    vis.clear();
    fix_(u, v, depth[v]+1);
    adj[u].push_back(v);
    adj[v].push_back(u);
    return 1;
  }
  template<typename F, typename G>
  void traverse_last_small(const F& f, const G& g) {
    vis.clear();
    traverse_(last_edge_.first, last_edge_.second, f, g);
  }
  template<typename F, typename G>
  void traverse_(int u, int p, const F& f, const G& g) {
    if(vis.get(u)) return;
    vis[u] = 1;
    f(u, p, adj[u]);
    for(int v : adj[u]) {
      if(v == p || vis.get(v)) continue;
      traverse_(v, u, f, g);
    }
    g(u, p, adj[u]);
  }
  void fix_(int u, int p, int d) {
    if(vis.get(u)) return;
    vis[u] = 1;
    pai[u] = p;
    depth[u] = d;
    for(int v : adj[u]) {
      if(v == p || vis.get(v)) continue;
      fix_(v, u, d+1);
    }
  }
  pair<int, int> last_edge() const {
    return last_edge_;
  }
  int parent(int i) const {
    return pai[i];
  }
};
} // namespace dsu
} // namespace lib



#line 5 "dsu/BinaryLifting.cpp"

namespace lib {
using namespace std;
namespace dsu {

template<typename D>
struct BinaryLifting : public D {
  using D::parent;
  vector<vector<int>> P;
  int K;

  BinaryLifting() : D() {}
  BinaryLifting(int n) : D(n) {
    P = decltype(P)(n, vector<int>(__lg(n)+1, -1));
    K = __lg(n)+1;
  }
  virtual void clear() override {
    D::clear();
    int n = P.size();
    P = decltype(P)(n, vector<int>(K, -1));
  }
  virtual int merge(int u, int v) override {
    if(!D::merge(u, v)) return 0;
    this->traverse_last_small([this](int u, int p, vector<int>&) {
      for(int& x : P[u]) x = -1;
      P[u][0] = p;
      for(int i = 1; i < K; i++) {
        if(P[u][i-1] == -1) break;
        P[u][i] = P[P[u][i-1]][i-1];
      }
    }, no_op_visitor);
    return 1;
  }
  int parent(int u, int k) {
    assert(k >= 0);
    for(int i = K-1; i >= 0; i--) {
      if(!((k>>i)&1)) continue;
      u = P[u][i];
      if(u == -1) return -1;
    }
    return u;
  }
};
} // namespace dsu
} // namespace lib

Back to top page