diff --git a/src/main.rs b/src/main.rs index e70a148..e4ddebd 100644 --- a/src/main.rs +++ b/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, 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, 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)); }