diff --git a/src/main.rs b/src/main.rs index 74c44e7..c3cdcbb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,20 +1,95 @@ -use std::ops::BitXor; +use std::cmp::{Ordering, min}; struct Solution; mod arr; - -impl Solution { - pub fn find_complement(num: i32) -> i32 { - if num == 0{ - 1 - }else { - num.bitxor((1 << (num.count_ones() + num.count_zeros() - num.leading_zeros())) - 1) +struct Dsu { + n: usize, + dt: Vec, + mb_cnt: Vec, + fulfilled: bool, +} +impl Dsu { + fn new(n: usize) -> Self { + let mut d = vec![0_usize; n]; + for i in 0..n { + d[i] = i; } + Self { + n, + dt: d, + mb_cnt: vec![1; n], + fulfilled: false, + } + } + fn get_father(&mut self, a: usize) -> usize { + if self.dt[a] == a { + return a; + } + self.dt[a] = self.get_father(self.dt[a]); + self.dt[a] + } + fn merge(&mut self, a: usize, b: usize) { + let af = self.get_father(a); + let bf = self.get_father(b); + self.dt[af] = bf; + self.mb_cnt[bf] += self.mb_cnt[af]; + if self.mb_cnt[bf] == self.n { + self.fulfilled = true; + } + } + fn is_fulfilled(&self) -> bool { + self.fulfilled + } + fn test_same(&mut self, a: usize, b: usize) -> bool { + self.get_father(a) == self.get_father(b) + } +} +impl Solution { + pub fn max_stability(n: i32, edges: Vec>, mut k: i32) -> i32 { + let mut heap = Vec::>::with_capacity(n as usize); + let mut dsu = Dsu::new(n as usize); + let mut min_val = 1e9 as i32; + for edge in edges { + if edge[3] == 1 { + if dsu.test_same(edge[0] as usize, edge[1] as usize) { + return -1; + } + dsu.merge(edge[0] as usize, edge[1] as usize); + min_val = min_val.min(edge[2]); + continue; + } + heap.push(edge); + } + heap.sort_unstable_by(|a, b| -> Ordering { b[2].cmp(&a[2]) }); + let mut used = Vec::::with_capacity((n / 2) as usize); + for i in 0..heap.len() { + let (a, b) = (heap[i][0] as usize, heap[i][1] as usize); + if dsu.test_same(a, b) { + continue; + } + dsu.merge(a, b); + used.push(heap[i][2]); + if dsu.is_fulfilled() { + for j in (0..used.len()).rev() { + if k > 0 { + min_val = min(min_val, used[j] * 2); + k -= 1; + } else { + min_val = min(min_val, used[j]); + break; + } + } + } + } + if dsu.fulfilled { min_val } else { -1 } } } fn main() { - let result = Solution::find_complement(0); - // let result = Solution::number_of_stable_arrays(200, 200, 25); + let result = Solution::max_stability( + 4, + vec![vec![0, 1, 3, 1], vec![1, 2, 3, 0], vec![2, 3, 1, 0]], + 1, + ); println!("{:?}", result); }