1962. Remove Stones to Minimize the Total

Priority Queue

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
    public int minStoneSum(int[] piles, int k) {
        Queue<Integer> maxHeap = new PriorityQueue<>((o1, o2) -> Integer.compare(o2, o1)); // 大顶堆
        for (int pile : piles) {
            maxHeap.offer(pile);
        }

        while (k > 0) {
            int pile = maxHeap.poll();
            maxHeap.offer(pile - pile / 2);
            k--;
        }

        int sum = 0;
        for (int pile : maxHeap) {
            sum += pile;
        }
        return sum;
    }
}

Heapify

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class Solution {
    public int minStoneSum(int[] piles, int k) {
        heapify(piles); // 原地堆化，保证堆顶节点为最大值

        while (k-- > 0) {
            piles[0] -= piles[0] / 2;
            sink(piles, 0); // 尝试将堆顶下潜
        }

        int sum = 0;
        for (int pile : piles) {
            sum += pile;
        }
        return sum;
    }

    private void heapify(int[] nums) {
        // 倒序遍历，保证 i 的左右子树已经堆化
        for (int i = nums.length / 2 - 1; i >= 0; i--) {
            sink(nums, i);
        }
    }

    private void sink(int[] nums, int i) {
        while (2 * i + 1 <= nums.length - 1) {
            // 存在子树，注意左节点的索引为 2 * i + 1
            int j = 2 * i + 1;
            if (j + 1 < nums.length && nums[j + 1] > nums[j]) {
                // 右节点值比左节点大
                j++;
            }

            if (nums[i] >= nums[j]) {
                // nums[i] 已经是最大值，无需交换
                break;
            }

            swap(nums, i, j);
            i = j; // 尝试继续下潜，把更大的值交换上来
        }
    }

    private void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}

References

1962. Remove Stones to Minimize the Total