3296. 移山所需的最少秒数

src/main.rs

@@ -1,95 +1,72 @@
-use std::cmp::{Ordering, min};
+use std::cmp::{max, Reverse};
+use std::collections::BinaryHeap;
 
 struct Solution;
 mod arr;
-struct Dsu {
-    n: usize,
-    dt: Vec<usize>,
-    mb_cnt: Vec<usize>,
-    fulfilled: bool,
+#[derive(Debug)]
+struct Node {
+    time: i32,
+    cnt: i32,
 }
-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,
-        }
+impl Node {
+    fn new(time: i32, cnt: i32) -> Node {
+        Node { time,cnt }
     }
-    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 total(&self) -> i64 {
+        let cnt = self.cnt as i64;
+        self.time as i64 * cnt * (cnt + 1) / 2
     }
-    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;
-        }
+}
+
+impl Ord for Node {
+    // 默认的比较方式
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.total().cmp(&other.total())
     }
-    fn is_fulfilled(&self) -> bool {
-        self.fulfilled
+}
+
+impl PartialOrd for Node {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.total().cmp(&other.total()))
     }
-    fn test_same(&mut self, a: usize, b: usize) -> bool {
-        self.get_father(a) == self.get_father(b)
+}
+
+impl Eq for Node { }
+
+impl PartialEq for Node {
+    fn eq(&self, other: &Self) -> bool {
+        self.total() == other.total()
+    }
+
+    fn ne(&self, other: &Self) -> bool {
+        self.total() != other.total()
     }
 }
 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);
+    pub fn min_number_of_seconds(mountain_height: i32, worker_times: Vec<i32>) -> i64 {
+        let mut heap = BinaryHeap::<Reverse<Node>>::new();
+        for x in worker_times {
+            heap.push(Reverse(Node::new(x, 1)));
         }
-        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;
-                    }
-                }
-            }
+        for _ in 0..mountain_height {
+            let mut node = heap.pop().unwrap();
+            node.0.cnt += 1;
+            heap.push(node);
         }
-        if dsu.fulfilled { min_val } else { -1 }
+        let mut ans = 0_i64; 
+        while !heap.is_empty() {
+            let mut node = heap.pop().unwrap().0;
+            node.cnt -= 1;
+            ans = max(ans, node.total())
+        }
+        ans
     }
 }
 
 fn main() {
-    let result = Solution::max_stability(
-        4,
-        vec![vec![0, 1, 3, 1], vec![1, 2, 3, 0], vec![2, 3, 1, 0]],
-        1,
+    let result = Solution::min_number_of_seconds(
+        5,
+        vec![1, 7],
     );
     println!("{:?}", result);
 }