Submission #60559108

---

Source Code Expand

#[allow(unused_imports)]
use itertools::Itertools;
use proconio::{input, marker::Usize1};

fn main() {
    input!{
        n: usize,
        m: usize,
        k: usize,
        e: [(Usize1, Usize1, u64); m],
        a: [Usize1; k],
        b: [Usize1; k],
    }
    let mut ans = 0;
    let f = |x: &(u64, u64), y: &(u64, u64)| {
        let a = x.0 + y.0;
        let b = x.1 + y.1;
        if a > b {
            (a - b, 0)
        } else {
            (0, b - a)
        }
    };

    let mut uf = UnionFind::new(n, f);

    let va = a.iter().fold(vec![0; n], |mut v, &ai| { v[ai]+=1; v });
    let vb = b.iter().fold(vec![0; n], |mut v, &ai| { v[ai]+=1; v });
    for i in 0..n {
        uf.insert_data(i, (va[i], vb[i]));
    }

    for &(u, v, w) in e.iter().sorted_unstable_by_key(|v| v.2) {
        if uf.same(u, v) {
            continue;
        }
        let (pa, pb) = uf.get_data(u).copied().unwrap();
        let (qa, qb) = uf.get_data(v).copied().unwrap();
        let val = (pa+qa).min(pb+qb);
        ans += val * w;
        uf.merge(u, v);
    }
    println!("{}", ans);
}

pub trait Debuggable {
    fn debug_string(&self) -> String;
}

impl<T: std::fmt::Debug + std::fmt::Display> Debuggable for Vec<T> {
    fn debug_string(&self) -> String {
        use itertools::Itertools;
        use std::iter::repeat;
        if let Some(max_size) = self.iter()
            .map(|x| format!("{:?}", x).len())
            .max() {
                let mut idx = String::from("idx |");   
                let mut val = String::from("val |");   
                for (i, xi) in self.iter().enumerate() {
                    idx.push_str(
                        &format!(" {:>w$} |", i, w=max_size)
                    );
                    val.push_str(
                        &format!(" {:>w$} |", xi, w=max_size)
                    );
                }

                format!("{}\n{}\n{}\n", idx, repeat('-').take(idx.len()).join(""), val)
            } else {
                format!("idx | \nval |\n")
            }
    }
}

impl<T: std::fmt::Debug + std::fmt::Display> Debuggable for std::collections::BTreeSet<T> {
    fn debug_string(&self) -> String {
        use itertools::Itertools;
        format!("{{ {} }}", self.iter().join(", "))
    }
}

impl<T, U> Debuggable for std::collections::BTreeMap<T, U> 
where T: std::fmt::Debug + std::fmt::Display, U: std::fmt::Debug + std::fmt::Display
{
    fn debug_string(&self) -> String {
        use itertools::Itertools;
        format!(
            "{{\n{}\n }}", self.iter()
                .map(|(k, v)| format!("{k} -> {v}, "))
                .join("\n")
        )
    }
}

// lg! マクロの定義
#[macro_export]
macro_rules! lg {
    ($val:expr) => {
        #[cfg(feature = "local")]
        {
            {
                use Debuggable;
                let val = &$val;
                eprintln!(
                    "[{}:{}] {} = \n{}",
                    file!(),
                    line!(),
                    stringify!($val),
                    val.debug_string()
                );
                val
            }
        }
    }
}
    #[derive(Debug)]
pub struct UnionFind<T, F>
where F: Fn(&T, &T) -> T,
{
    vertex: Vec<usize>,
    data: Vec<Option<T>>,
    merge_op: F,
}

impl<T: Clone, F: Fn(&T, &T) -> T> UnionFind<T, F> {
    pub fn new(size: usize, merge_op: F) -> Self {
        UnionFind {
            vertex: vec![!1; size] ,
            data: vec![None; size],
            merge_op,
        }
    }

    pub fn leader(&mut self, u: usize) -> usize {
        let elm: usize = self.vertex[u];
        if elm > self.vertex.len() {
            // 親ノードなので返却
            u
        } else {
            // 子ノードなので親ノードを取得したのち、圧縮
            self.vertex[u] = self.leader(elm);
            self.vertex[u]
        }
    }

    pub fn same(&mut self, u: usize, v: usize) -> bool {
        self.leader(u) == self.leader(v)
    }

    pub fn size(&mut self, u: usize) -> usize {
        let idx: usize = self.leader(u);
        !self.vertex[idx]
    }

    pub fn merge(&mut self, u: usize, v: usize) -> bool {
        if self.same(u, v) {
            // すでに親が同じ場合は何もせず、「マージされなかった」ことを報告
            return false;
        }
        let u_leader = self.leader(u);
        let u_size = self.size(u);

        let v_leader = self.leader(v);
        let v_size = self.size(v);

        let merged_size = !(u_size + v_size);
        if u_size < v_size {
            self.vertex[v_leader] = merged_size;
            self.vertex[u_leader] = v_leader;

            self.data[v_leader] = match (&self.data[u_leader], &self.data[v_leader]) {
                (Some(du), Some(dv)) => Some((self.merge_op)(du, dv)),
                (None, None) => None,
                _ => { unreachable!(); }
            };
        } else {
            self.vertex[u_leader] = merged_size;
            self.vertex[v_leader] = u_leader;

            self.data[u_leader] = match (&self.data[u_leader], &self.data[v_leader]) {
                (Some(du), Some(dv)) => Some((self.merge_op)(du, dv)),
                (None, None) => None,
                _ => { unreachable!(); }
            };
        }

        true
    }

    pub fn insert_data(&mut self, u: usize, value: T) {
        let root = self.leader(u);
        self.data[root] = Some(value);
    }

    pub fn get_data(&mut self, u: usize) -> Option<&T> {
        let root = self.leader(u);
        self.data[root].as_ref()
    }

}

Submission Info

Compile Error

warning: the item `Itertools` is imported redundantly
  --> src/main.rs:52:13
   |
2  | use itertools::Itertools;
   |     -------------------- the item `Itertools` is already imported here
...
52 |         use itertools::Itertools;
   |             ^^^^^^^^^^^^^^^^^^^^
   |
   = note: `#[warn(unused_imports)]` on by default

warning: the item `Itertools` is imported redundantly
  --> src/main.rs:77:13
   |
2  | use itertools::Itertools;
   |     -------------------- the item `Itertools` is already imported here
...
77 |         use itertools::Itertools;
   |             ^^^^^^^^^^^^^^^^^^^^

warning: the item `Itertools` is imported redundantly
  --> src/main.rs:86:13
   |
2  | use itertools::Itertools;
   |     -------------------- the item `Itertools` is already imported here
...
86 |         use itertools::Itertools;
   |             ^^^^^^^^^^^^^^^^^^^^

Judge Result