각 건물 쌍은 3가지 중 하나의 관계를 맺고 있습니다.
- 서로 같은 나무를 심어야 한다.
- 서로 다른 나무를 심어야 한다.
- 제약이 없다.
단순하게 제약이 있는 나무들을 하나의 집합으로 만들어 관리하면 서로소 집합 자료구조를 이용할 수 있습니다.
단, 집합의 루트를 찾을 때 루트-리프 사이 관계를 계속 갱신해줘야 하는데, 서로 같은 나무를 심는 조건을 0, 서로 다른 나무를 심는 조건을 1로 두면 xor 연산을 통해 조건의 합성을 정의할 수 있습니다.
경로 압축 시에 조건을 합성해주는 코드를 추가하여 아래 그림처럼 만들 수 있습니다.
루트에 각각 은행나무, 플라타너스를 심었을 때 집합에서 얻을 수 있는 최솟값을 저장하고 질의마다 갱신할 수 있습니다.
코드
use std::fmt::Write;
use std::io::{stdin, Read};
fn main() {
let mut input = String::new();
stdin().read_to_string(&mut input).unwrap();
let mut input = input.split_ascii_whitespace().map(str::parse).flatten();
let n = input.next().unwrap();
let d = input.next().unwrap();
let mut costs = Vec::with_capacity(n);
for _ in 0..n {
let gi = input.next().unwrap();
let pi = input.next().unwrap();
costs.push((gi, pi));
}
let mut forest = SetForest::new(n, costs.clone());
for _ in 0..d {
let c = input.next().unwrap();
let i = input.next().unwrap() - 1;
let j = input.next().unwrap() - 1;
forest.merge(i, j, c == 1);
}
let mut output = String::new();
writeln!(output, "{}", forest.min).unwrap();
let q = input.next().unwrap();
for _ in 0..q {
let c = input.next().unwrap();
let a = input.next().unwrap() - 1;
let b = input.next().unwrap();
match c {
0 => forest.merge(a, b - 1, false),
1 => forest.merge(a, b - 1, true),
2 => {
forest.update(a, costs[a], (b, costs[a].1));
costs[a].0 = b;
}
_ => {
forest.update(a, costs[a], (costs[a].0, b));
costs[a].1 = b;
}
}
writeln!(output, "{}", forest.min).unwrap();
}
print!("{}", output);
}
struct SetForest {
parent: Vec<usize>,
rank: Vec<usize>,
conflicts: Vec<bool>,
costs: Vec<(usize, usize)>,
min: usize,
}
impl SetForest {
fn new(s: usize, costs: Vec<(usize, usize)>) -> Self {
let min = costs.iter().map(|(a, b)| a.min(b)).sum();
Self {
parent: (0..s).collect(),
rank: vec![1; s],
conflicts: vec![false; s],
costs,
min,
}
}
fn find(&mut self, i: usize) -> (usize, bool) {
if i == self.parent[i] {
(i, false)
} else {
let (root, conflicts) = self.find(self.parent[i]);
self.parent[i] = root;
self.conflicts[i] = conflicts != self.conflicts[i];
(root, self.conflicts[i])
}
}
fn update(&mut self, i: usize, (pa, pb): (usize, usize), (ua, ub): (usize, usize)) {
let (pi, ci) = self.find(i);
let (pia, pib) = self.costs[pi];
let (newa, newb) = if ci {
(pia - pb + ub, pib - pa + ua)
} else {
(pia - pa + ua, pib - pb + ub)
};
self.costs[pi] = (newa, newb);
self.min -= pia.min(pib);
self.min += newa.min(newb);
}
fn merge(&mut self, u: usize, v: usize, conflicts: bool) {
let (mut pu, cu) = self.find(u);
let (mut pv, cv) = self.find(v);
if pu == pv {
return;
}
if self.rank[pu] < self.rank[pv] {
std::mem::swap(&mut pu, &mut pv);
}
self.parent[pv] = pu;
if self.rank[pu] == self.rank[pv] {
self.rank[pu] += 1;
}
self.conflicts[pv] = (cu != cv) != conflicts;
let (ua, ub) = self.costs[pu];
let (va, vb) = self.costs[pv];
self.min -= ua.min(ub);
self.min -= va.min(vb);
let (newa, newb) = if self.conflicts[pv] {
(ua + vb, ub + va)
} else {
(ua + va, ub + vb)
};
self.costs[pu] = (newa, newb);
self.min += newa.min(newb);
}
}