ぺんぎんメモ

プログラミングのメモです。たまに私生活のことや鬱っぽいことを書きます。

最小限のHL分解の実装

色々機能を追加する前の実装。

参考にしたページ
http://beet-aizu.hatenablog.com/entry/2017/12/12/235950
https://codeforces.com/blog/entry/53170

struct HLD {
  vector<vector<int>> g;
  vector<int> vid;
  vector<int> head;
  HLD(vector<vector<int>> g) : g(g) {
    int n = g.size();
    vid = vector<int>(n);
    head = vector<int>(n);
    {
      vector<int> sz(n);
      auto dfs = [&](auto dfs, int u) -> void {
        sz[u] = 1;
        for (auto &v : g[u]) {
          dfs(dfs, v);
          sz[u] += sz[v];
          if (sz[v] > sz[g[u][0]]) swap(v, g[u][0]);
        }
      };
      dfs(dfs, 0);
    }
    {
      vector<int> next(n, -1);
      queue<int> q;
      auto dfs = [&](auto dfs, int u) -> void {
        for (auto v : g[u]) {
          if (v == g[u][0]) {
            next[u] = v;
            head[v] = head[u];
          } else {
            q.emplace(v);
            head[v] = v;
          }
          dfs(dfs, v);
        }
      };
      auto bfs = [&](int s) {
        int k = 0;
        while (q.size()) {
          int u = q.front(); q.pop();
          for (int v = u; v != -1; v = next[v]) {
            vid[v] = k++;
          }
        }
      };
      q.emplace(0);
      dfs(dfs, 0);
      bfs(0);
    }
  }
};