3346. 执行操作后元素的最高频率 I

src/main.rs

@@ -1,26 +1,29 @@
+use std::cmp::{max, min};
 mod arr;
 struct Solution {}
 
 impl Solution {
-    pub fn maximum_energy(mut energy: Vec<i32>, k: i32) -> i32 {
-        let mut ans = -1e9 as i32;
-        for i in 1..energy.len() {
-            if i as i32 - k >= 0 {
-                if energy[i - k as usize] > 0 {
-                    energy[i] += energy[i - k as usize];
-                }
-            }
-            if i as i32 + k >= energy.len() as i32 {
-                ans = ans.max(energy[i]);
-            }
+    pub fn max_frequency(nums: Vec<i32>, k: i32, num_operations: i32) -> i32 {
+        let mut ans = 0;
+        let mut arr = vec![0; 1_000_00 + 10];
+        let mut arr2 = vec![0; 1_000_00 + 10];
+        for &num in nums.iter() {
+            let l = max(num - k, 0);
+            let r = min(1_000_00 + 1, num + k);
+            arr[l as usize] += 1;
+            arr[r as usize + 1] -= 1;
+            arr2[num as usize] += 1;
+        }
+        let mut base = 0;
+        for i in 0..arr.len() {
+            base += arr[i];
+            let local_ans = min(arr2[i] + num_operations, base);
+            ans = max(ans, local_ans);
         }
         ans
     }
 }
 
 fn main() {
-    // [[1,2,2,3,5],[3,2,3,4,4],[2,4,5,3,1],[6,7,1,4,5],[5,1,1,2,4]]
-    // initial as a 2D vector
-    let skills = vec![5, 2, -10, -5, 1];
-    println!("{:?}", Solution::maximum_energy(skills, 3));
+    println!("{:?}", Solution::max_frequency(vec![1, 1, 1, 4, 5], 1, 2));
 }