bouzuya.hatenablog.com

ぼうずやのにっき

とても寒い / ABC328 F

とても寒い。暖房をつけている。毛布も出している。寒さと毛布のホコリで鼻水が止まらない。

とりあえず Bluesky に登録してみた。 https://bsky.app/profile/bouzuya.bsky.social 。

ABC328 の F https://atcoder.jp/contests/abc328/tasks/abc328_f 。

提出 https://atcoder.jp/contests/abc328/submissions/47540177 。

解説 AC 。重み付き Union-Find 。 Union-Find の応用。きちんと Union-Find を理解していれば解けるのかもしれない……。次に類似の問題が出たら解けるようにしたい。

use proconio::{input, marker::Usize1};

// 重み付きUnionFind
struct WeightedUnionFind {
    parent: Vec<usize>,
    weight: Vec<i64>,
}

impl WeightedUnionFind {
    fn new(n: usize) -> Self {
        Self {
            parent: (0..n).collect(),
            weight: vec![0; n],
        }
    }

    fn merge(&mut self, a: usize, b: usize, w: i64) -> bool {
        let (ra, rb) = (self.root(a), self.root(b));
        if ra == rb {
            return self.weight[b] - self.weight[a] == w;
        }
        self.parent[ra] = rb;
        self.weight[ra] = self.weight[b] - self.weight[a] - w;
        true
    }

    fn root(&mut self, a: usize) -> usize {
        if self.parent[a] == a {
            return a;
        }
        let p = self.parent[a];
        let r = self.root(p);
        self.parent[a] = r;
        self.weight[a] += self.weight[p];
        r
    }
}

fn main() {
    input! {
        _n: usize,
        q: usize,
        abd: [(Usize1, Usize1, i64); q],
    };

    let mut dsu = WeightedUnionFind::new(2 * 100_000_usize + 1);
    for (i, (a, b, d)) in abd.into_iter().enumerate() {
        if dsu.merge(a, b, d) {
            println!("{}", i + 1);
        }
    }
}

今日のコミット。