백준 4014 - 도로

자갈도로 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
    }
}