Rustコトハジメ

プログラミング言語Rustと競プロに関する情報をお届けします。

【ARC065-B 連結】Union Findを2つ使う

UFを2つ使うっぽいことはわかるが、いざ書こうと思うと書けなかった。結局わかってなかったという話。

D - 連結 / Connectivity

問題

N個の町がある。そこにK本の道路とL本の線路がある。

町iと町jは道路を使っても行き来出来るし、線路を使っても行き来出来る

を満たすjの数を、各iについて求めなさい。

方針

線路がなければUFで終わりの問題。

しかし、道路でも線路でも行き来出来る <=> 道路についてのUFでも線路についてのUFでも同じグループにいる

ということに気づけばおわり。(道路グループ、線路グループ)のペアをキーにしてハッシュマップで実装。

実装

fn solve() {
    input!{
        N:usize,K:usize,L:usize,
        PQ:[(usize,usize);K],
        RS:[(usize,usize);L]
    }
    let mut roadUF=UnionFind::new(N);
    for (p,q) in PQ {
        roadUF.merge(p-1, q-1);
    }
    let mut railUF=UnionFind::new(N);
    for (r,s) in RS {
        railUF.merge(r-1, s-1);
    }
 
    // 2つのi,jについて、
    // (roadUF[i],railUF[i]) == (roadUF[j],railUF[j])
    // の時、2つのノードは道路と鉄道の両方で接続されている
    let mut ty = vec![];
    let mut h = HashMap::new();
    for i in 0..N {
        let t = (roadUF.root(i), railUF.root(i));
        ty.push(t);
        *h.entry(t).or_insert(0) += 1;
    }
    let mut ans = vec![];
    for i in 0..N {
        let t = ty[i];
        ans.push(*h.get(&t).unwrap());
    }
 
    println!("{}", ans.iter().map(|&x| x.to_string()).collect::<Vec<String>>().connect(" "));
}

考察

最近のABCはUFあまり出てないけど昔の問題を解くと時々出てくる。

他の問題としては、

  • D - Decayed Bridges: 橋が落ちていくから、その時々の不便度を計算せよという問題。UFを使って橋を逆に繋いでいって不便度を都度更新していく
  • D - Equals: スワップをグラフとして考えて、いくつかの非連結な成分に分解し、その中で個数を数える。この時、UFが使える

などがありそう。