자갈도로 K개를 포함하는 스패닝 트리를 찾는 문제입니다.
불가능한 경우를 일단 거르기 위해 자갈 도로를 최대한 사용해서 만들 수 있는 트리를 살펴봤습니다.
이 트리의 간선이 K개 미만이면 K개의 무료 자갈 도로를 사용했을 때 오직 하나의 무료도로로 이루어진 경로가 있다는 조건에 위배됩니다.
또한 이 트리의 정보를 최대한 이용할 수 없을까 생각하다가 아이디어를 떠올렸습니다.
어떤 콘크리트 도로가 이 트리에서 이미 이어진 두 정점을 잇는다면, 임의의 자갈 도로를 제거하고 이 콘크리트 도로를 추가하는 것으로 정점 간 연결성을 유지하되 자갈 도로의 개수를 하나 줄일 수 있다는 것입니다.
이 작업을 가능한 모든 콘크리트 도로에 대해 반복하면 K개의 무료 자갈 도로를 사용한 최대한의 트리를 만들 수 있습니다.
이 트리가 모든 정점을 잇는다면 도로유지방안이 존재하고, 아니면 없습니다.
각 단계에서 모든 간선을 확인하고 연결성을 확인하기 때문에 시간은 O(Nlog*N)입니다.
코드
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 output = String::new();
let mut input = input.split_ascii_whitespace().map(str::parse).flatten();
let n = input.next().unwrap();
let m = input.next().unwrap();
let k = input.next().unwrap();
let mut concretes = vec![];
let mut gravels = vec![];
for _ in 0..m {
let ui = input.next().unwrap();
let vi = input.next().unwrap();
let ci = input.next().unwrap();
if ci == 0 {
gravels.push((ui, vi));
} else {
concretes.push((ui, vi));
}
}
let mut gravel_paths = Sets::new(n + 1);
let mut components = n;
let mut gravel_count = 0;
for &(ui, vi) in &gravels {
let ru = gravel_paths.root(ui);
let rv = gravel_paths.root(vi);
if ru != rv {
gravel_paths.join(ru, rv);
components -= 1;
gravel_count += 1;
}
}
if gravel_count < k {
println!("no solution");
} else {
let mut paths = Sets::new(n + 1);
let mut list = Vec::with_capacity(n);
for &(ui, vi) in &concretes {
let gu = gravel_paths.root(ui);
let gv = gravel_paths.root(vi);
let ru = paths.root(ui);
let rv = paths.root(vi);
if ru != rv {
if gu != gv {
paths.join(ru, rv);
gravel_paths.join(gu, gv);
list.push((ui, vi, 1));
components -= 1;
} else if gravel_count > k {
paths.join(ru, rv);
list.push((ui, vi, 1));
gravel_count -= 1;
}
}
}
if components != 1 || gravel_count > k {
println!("no solution");
} else {
for &(ui, vi) in &gravels {
let ru = paths.root(ui);
let rv = paths.root(vi);
if ru != rv {
paths.join(ru, rv);
list.push((ui, vi, 0));
}
}
for (ui, vi, ci) in list {
writeln!(output, "{} {} {}", ui, vi, ci).ok();
}
print!("{}", output);
}
}
}
struct Sets {
up: Vec<usize>,
rank: Vec<usize>,
}
impl Sets {
fn new(n: usize) -> Self {
Self {
up: (0..n).collect(),
rank: vec![1; n],
}
}
fn root(&mut self, mut x: usize) -> usize {
while x != self.up[x] {
self.up[x] = self.up[self.up[x]];
x = self.up[x];
}
x
}
fn join(&mut self, mut rx: usize, mut ry: usize) -> usize {
if self.rank[rx] < self.rank[ry] {
std::mem::swap(&mut rx, &mut ry);
}
self.up[ry] = rx;
if self.rank[rx] == self.rank[ry] {
self.rank[rx] += 1;
}
rx
}
}