LeetCode Entry

3075. Maximize Happiness of Selected Children

09.05.2024 medium 2024 kotlin rust

Sum of k maximums decreasing each step

3075. Maximize Happiness of Selected Children medium blog post substack youtube 2024-05-09_11-24.webp

Join me on Telegram

https://t.me/leetcode_daily_unstoppable/597

Problem TLDR

Sum of k maximums decreasing each step #medium #sorting #heap #quickselect

Intuition

By the problem definition we may assume that optimal solution is to take the largest values first, as smaller values will not decrease the result after reaching zero.

There are several ways to take k largest values: sort the entire array, use Heap (PriorityQueue) or use QuickSelect and sort partially.

Approach

Let’s use PriorityQueue in Kotlin (min heap) and QuickSelect in Rust (select_nth_unstable).

  • when using heap we can take at most k values into it to save space and time
  • Rust’s select_nth_unstable result tuple is not very easy to use (do you know a better way?)

Complexity

  • Time complexity: \(O(n + klog(k))\) for the Heap and for the QuickSelect

  • Space complexity: \(O(n)\) for the Heap, \(O(1)\) for the QuickSelect

Code


    fun maximumHappinessSum(happiness: IntArray, k: Int): Long {
        val pq = PriorityQueue<Int>()
        for (h in happiness) { pq += h; if (pq.size > k) pq.poll() }
        return (0..<k).sumOf { max(0, pq.poll() + it - k + 1).toLong() }
    }



    pub fn maximum_happiness_sum(mut happiness: Vec<i32>, k: i32) -> i64 {
        let count = 0.max(happiness.len() as i32 - k - 1) as usize;
        let gt = if count > 0 { happiness.select_nth_unstable(count).2 }
                 else { &mut happiness[..] };
        gt.sort_unstable_by(|a, b| b.cmp(a));
        (0..k).map(|i| 0.max(gt[i as usize] - i) as i64).sum()
    }