ぺんぎんメモ

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

Tree雑記②

LCA

LCAが思っていた以上に使えることがわかった。
重み付きの木で「ある頂点からある頂点までのコストを求めよ」という105個くらいのクエリに答えたいとき、別に全方位木DPを使わなくてもよくて、LCAで事足りる。
LCAの準備の過程で根からの深さも求まっているので、これを利用して2頂点間のパスも効率的に求められる(LCAにたどり着くまで親方向に進む)。

AtCoderの解説pdfに「木で累積和を取ってLCA」みたいな文言が書かれていて、そのときはよくわからなかったけど、今なら少しわかる。dist[u] + dist[v] - dist[lca(u, v)] * 2みたいなことをすれば、u - v間の何か(距離であったりコストであったり)を効率的に求められる。

ダブリング

久しぶりにこちらの問題を解いた。この問題はダブリングで解くのだけど、当時はダブリングへの理解が曖昧な状態で実装していたため、微妙に腑に落ちなかった。しかし、昨日実装してみてようやく腑に落ちた。

ということで、LCAのプログラムをソラで書けるようになった。書いたコードは以下(nは頂点数、gvector<vector<Edge>>型、その他はこちらのテンプレートを参照)。

vector<vi32> par(30, vi32(n + 1, n));
vi32 dep(n);
{
  function<void(int, int, int)> dfs = [&](int u, int p, int d) {
    if (p != -1) par[0][u] = p;
    dep[u] = d;
    each(e, g[u]) {
      if (e.to == p) continue;
      dfs(e.to, u, d + 1);
    }
  };
  dfs(0, -1, 0);
}
rep(i, 29) rep(j, n) {
  par[i + 1][j] = par[i][par[i][j]];
}
function<int(int, int)> lca = [&](int u, int v) {
  if (dep[u] < dep[v]) swap(u, v);
  rep(i, 30) {
    if ((dep[u] - dep[v]) & bit32(i)) u = par[i][u];
  }
  if (u == v) return u;
  rrep(i, 30) {
    if (par[i][u] != par[i][v]) {
      u = par[i][u], v = par[i][v];
    }
  }
  return par[0][u];
};

par[i][par[i][j]]で配列外参照を起こさないように、子孫が存在しない場合にnを持つようにしている。その影響で、parの二次元配列の2つ目の添え字の大きさもn + 1になっている(このあたりは、すぬけさんの実装を参考にしている)。

par[0][u] = p;if (e.to == p) continue;の後に書こうか悩んでいる。こうすることで、if (p != -1)という条件分岐が不要になる。

function<void(int, int, int)> dfs = [&](int u, int p, int d) {
  dep[u] = d;
  each(e, g[u]) {
    if (e.to == p) continue;
    par[0][e.to] = u;
    dfs(e.to, u, d + 1);
  }
};

こちらのほうがいいかも。

というか、if (e.to == p) continue;というコードにわざわざ一行使いたくない。

function<void(int, int, int)> dfs = [&](int u, int p, int d) {
  dep[u] = d;
  each(e, g[u]) if (e.to != p) {
    par[0][e.to] = u;
    dfs(e.to, u, d + 1);
  }
};

dep[u] = d;も移動させたくなった。ついでに引数のdもなくした。

function<void(int, int)> dfs = [&](int u, int p) {
  each(e, g[u]) if (e.to != p) {
    par[0][e.to] = u;
    dep[e.to] = dep[u] + 1;
    dfs(e.to, u);
  }
};

コードを最適化していく過程で、あることに気付いた。each内であれば辺の情報も取得できるのに対し、eachの外だと辺の情報が取得しづらい。

function<void(int, int)> dfs = [&](int u, int p) {
  /* hard to get edge information */
  each(e, g[u]) if (e.to != p) {
    /* easy to get edge information */
    dfs(e.to, u);
  }
};

eachの中のほうが色々と楽。あと、eachの中だと一度のDFSで色々なことを同時にしやすい。たとえば、次のコードは3つのことを同時に行っている。

function<void(int, int)> dfs = [&](int u, int p) {
  each(e, g[u]) if (e.to != p) {
    par[0][e.to] = u;
    dep[e.to] = dep[u] + 1;
    dist[e.to] = dist[u] + e.cost;
    dfs(e.to, u);
  }
};

これからはこの書き方をする。