diff --git a/src/problem/mod.rs b/src/problem/mod.rs index 43857e8..23ce6f1 100644 --- a/src/problem/mod.rs +++ b/src/problem/mod.rs @@ -67,4 +67,5 @@ mod p232_implement_queue_using_stacks; mod p1976_number_of_ways_to_arrive_at_destination; mod p2917_find_the_k_or_of_an_array; mod p2575_find_the_divisibility_array_of_a_string; -mod p2834_find_the_minimum_possible_sum_of_a_beautiful_array; \ No newline at end of file +mod p2834_find_the_minimum_possible_sum_of_a_beautiful_array; +mod p2386_find_the_k_sum_of_an_array; \ No newline at end of file diff --git a/src/problem/p2386_find_the_k_sum_of_an_array.rs b/src/problem/p2386_find_the_k_sum_of_an_array.rs new file mode 100644 index 0000000..5259168 --- /dev/null +++ b/src/problem/p2386_find_the_k_sum_of_an_array.rs @@ -0,0 +1,84 @@ +/** + * [2386] Find the K-Sum of an Array + */ +pub struct Solution {} + +// submission codes start here +use std::collections::BinaryHeap; + +#[derive(PartialEq, Eq)] +struct Node { + sum: i64, + last: usize, +} + +impl PartialOrd for Node { + fn partial_cmp(&self, other: &Self) -> Option { + other.sum.partial_cmp(&self.sum) + } +} + +impl Ord for Node { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + other.sum.cmp(&self.sum) + } +} + +impl Solution { + pub fn k_sum(nums: Vec, k: i32) -> i64 { + let mut total = 0; + let mut nums: Vec = nums + .iter() + .map(|x| { + if *x > 0 { + let x = *x as i64; + total += x; + x + } else { + -x as i64 + } + }) + .collect(); + + nums.sort_unstable(); + + let mut heap = BinaryHeap::new(); + heap.push(Node { + sum: nums[0], + last: 0, + }); + + let mut result = 0; + for _ in 1..k { + let top = heap.pop().unwrap(); + + result = top.sum; + if top.last == nums.len() - 1 { + continue; + } + + heap.push(Node { + sum: top.sum + nums[top.last + 1], + last: top.last + 1, + }); + heap.push(Node { + sum: top.sum - nums[top.last] + nums[top.last + 1], + last: top.last + 1, + }); + } + + total - result + } +} + +// submission codes end + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_2386() { + assert_eq!(2, Solution::k_sum(vec![2, 4, -2], 5)); + } +}