マージテクで、いわゆる auxiliary tree を構築する方法

参考問題：

ABC340-G Leaf Color （鹿島建設プログラミングコンテスト2024）

参考文献：

よく知られたアルゴリズムでは、まず任意の $2$ 頂点の間の LCA を計算できる状態にして、それから各色の頂点を行きがけ順などでソートして走査していくと思うのですが、LCA は一から実装すると意外と大変。そこでゼロスタートでも簡単に書ける方法を考えました。（というか知らなくて再発明しました。）計算量は変わらず $O ( N \log N)$ です。

入力

g：子リスト
ord: 木の頂点を根から順に並べたリスト（たとえば行きがけ順）
color: 各頂点の色

途中に出てくる変数

lca[x][&c]: usize: x の部分木の色 c の節点全体の LCA（ただしマージテクで計算しているので毎回破壊されています）
h[x][&c]: Vec<usize>：色 c に対応するいわゆる auxiliary tree における x の子リスト
head[&c]: usize：lca[y][&c] たちのうちのどれでもいいのでひとつ（このコードでは一番最後に登場するやつ）
tail[&c]: Vec<usize>：lca[y][&c] たちのうち head[c] でないものすべて

方針

(1) lca[x][&c] をマージテクで計算します。初手で lca[y][&c] たちを head[&c] と tail[&c] に分類しているのですが、これは $(A, B) \rightarrow (A \cap B, A\cup B)$ みたいな破壊的な二項演算を畳み込んでいると思うとわかりやすいと思います。集合としての union が head[&c] に集められて行って、あぶれた要素が tail[&c] に再就職していく流れです。（登場順でいうと head に入っているやつが一番最後なの命名ひどくて申し訳ないです。）

(2) 大本命の h[x][&c] です。いわゆる auxiliary tree に x が登場するのは、color[x] == x のときか、lca[y][&c] が定義されている y が複数存在するとじゅですね。前者は tail[x] をすべて見れば必要な色がすべてそろいます。子リストも tail[x] を見ればだいたいそろいますけれども、ひとつだけ head[x] のほうに入っているので忘れずに回収しましょう。後者はみるべき色が1つなので楽です。前者と重複して挿入しないように注意しましょう。

ソースコード

これを使った提出👉 Submission #51214479 - KAJIMA CORPORATION CONTEST 2024（AtCoder Beginner Contest 340）

Submission #51214547 - KAJIMA CORPORATION CONTEST 2024（AtCoder Beginner Contest 340）

fn auxialiary_tree(
    g: &[Vec<usize>],
    ord: &[usize],
    color: &[usize],
) -> Vec<HashMap<usize, Vec<usize>>> {
    let mut lca = vec![HashMap::new(); g.len()];
    let mut h = vec![HashMap::<_, Vec<_>>::new(); g.len()];
    for &x in ord.iter().rev() {
        // (1 前半) `head` と `tail` の計算
        let mut head = HashMap::new();
        let mut tail = HashMap::new();
        for &y in &g[x] {
            if head.len() < lca[y].len() {
                swap(&mut head, &mut lca[y]);
            }
            for (&c, &z) in lca[y].iter() {
                if let Some(z) = head.insert(c, z) {
                    tail.entry(c).or_insert_with(Vec::new).push(z);
                }
            }
        }
        // (2 の場合分け1つ目) `color[x] == x` のとき、かつ場合分け1つ目に該当しないとき
        if !tail.contains_key(&color[x]) {
            h[x].entry(color[x]).or_default();
            if let Some(&y) = head.get(&color[x]) {
                h[x].get_mut(&color[x]).unwrap().push(y);
            }
        }
        // (2 の場合分け2つ目) `lca[y][&c]` が定義されている `y` が複数存在するとき
        for (&c, tail) in &tail {
            h[x].insert(c, vec![head[&c]]);
            h[x].entry(c).or_default().extend(tail);
        }
        // (1 後半) `lca[x]`  の計算
        lca[x] = head;
        for &c in tail.keys() {
            *lca[x].get_mut(&c).unwrap() = x;
        }
        lca[x].insert(color[x], x);
    }
    h
}

ブログ名

競技プログラミングやお歌のお話をする高菜です。

マージテクで、いわゆる auxiliary tree を構築する方法

入力

途中に出てくる変数

方針

ソースコード