3600. 升级后最大生成树稳定性

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<usize>,
+    mb_cnt: Vec<usize>,
+    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<Vec<i32>>, mut k: i32) -> i32 {
+        let mut heap = Vec::<Vec<i32>>::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::<i32>::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);
 }