Leetcode 2602:前缀和+二分查找

文章内容

一、题意

题目连接:使数组元素全部相等的最少操作次数

给你一个数组 nums,再给一堆查询 queries。对每个查询值 q,你可以做任意次操作,每次挑数组里的某个元素,给它 +1 或者 -1。问:最少多少次操作,能让数组里所有元素都变成 q?

每个 query 都要单独回答一次,互不影响
举个例子,nums = [3, 1, 6, 8],q = 1。3 要变成 1,减 2 次;1 不用动;6 要减 5 次;8 要减 7 次。加起来 14 次。

把所有数组元素统一到目标值的操作概念图

暴力解法:对每个 query 遍历一遍数组,把 |nums[i] - q| 全加起来。但是 nums 和 queries 都能到 10^5 这个量级,O(n × m) 就是 10^10 次运算会导致超时。

二、二分模板

lower_bound二分模板代码:在一个有序数组里,找到第一个大于等于目标值的位置

java
private int lowerBinary(int[] nums, int checkValue, int l, int r) {    while (l < r) {        int mid = l + (r - l) / 2;        if (nums[mid] >= checkValue) {            r = mid;                  } else {            l = mid + 1;      // mid 太小了,答案在右边,可以跳过 mid        }    }    return l;}

三、解题思路

第一步,排序。 排完序之后,整个数组就有了一个性质:任意一个 q 值,都能把数组干净利落地劈成两段。左边全是小于 q 的,右边全是大于等于 q 的,中间不会交错。

二分查找将有序数组分成左右两部分

左边那些数要加到 q,右边那些数要减到 q。这就把绝对值拆掉了——不用再一个个判断谁大谁小。

第二步,前缀和。 建一个 prefixSum 数组,prefixSum[i] 表示前 i 个元素的和。有了它,求任意一段区间的和就是一次减法,O(1)。这是整道题能跑起来的关键:我们要反复问"左边这堆数加起来多少",不可能每次现算。

第三步,对每个 query 二分找分界点。 分界点就是 lower_bound 返回的那个索引,记作 midIndex。

过程图解:

前缀和公式推导:左右两边的操作次数计算

左边有 midIndex 个元素(索引从 0 开始,所以第一个 >= q 的位置的下标,正好等于它前面元素的个数),每个都要涨到 q。如果它们本来的和是 S_left,那总共要加的量就是 q * midIndex - S_left。

右边有 n - midIndex 个元素,本来的和是 S_right,每个都要降到 q,总共要减的量是 S_right - q * (n - midIndex)。

两个一加,就是这个 query 的答案。整个过程只用了一次二分和几次加减乘,O(log n)。

四、代码

java
class Solution {    /**     * 二分查找 lower_bound:在 [l, r) 区间内找到第一个 >= checkValue 的位置。     * 如果不存在这样的元素,返回 r。     */    private int lowerBinary(int[] nums, int checkValue, int l, int r) {        while (l < r) {            int mid = l + (r - l) / 2;            if (nums[mid] >= checkValue) {                r = mid;        // mid 可能是答案,保留            } else {                l = mid + 1;    // mid 一定不是答案,跳过            }        }        return l;    }     public List<Long> minOperations(int[] nums, int[] queries) {        // 排序是整个解法的前提,不排序没法二分        Arrays.sort(nums);         // 前缀和数组,长度是 n+1,prefixSum[i] = nums[0] + ... + nums[i-1]        // 多出来的那一位(prefixSum[0] = 0)省掉了边界判断        long[] prefixSum = new long[nums.length + 1];        for (int i = 1; i <= nums.length; i++) {            prefixSum[i] = prefixSum[i - 1] + nums[i - 1];        }         List<Long> res = new ArrayList<>();        for (int query : queries) {            // 分界点:第一个 >= query 的位置            // 它同时也是"左边有多少个元素"            int midIndex = lowerBinary(nums, query, 0, nums.length);             // 左半段:所有 < query 的元素,每个都要涨到 query            long ans = (long) query * midIndex - prefixSum[midIndex];             // 右半段:所有 >= query 的元素,每个都要降到 query            // prefixSum[prefixSum.length - 1] 就是全部元素之和            ans += (prefixSum[prefixSum.length - 1] - prefixSum[midIndex])                 - (long) query * (nums.length - midIndex);             res.add(ans);        }        return res;    }}

前缀和数组长度取 n + 1 是个小技巧。prefixSum[0] = 0 作为哨兵,可以让 midIndex = 0(所有元素都 >= query)这种情况不需要特判。

五、复杂度

排序 O(n log n),建前缀和 O(n),每个 query 一次二分 O(log n),一共 m 个 query。

总的就是 O(n log n + m log n)。跟暴力的 O(n × m) 比,从 10^10 降到了大概 3 × 10^6,差了三个数量级。


本文由 AgentPlanFlow 生成

觉得不错?给作者点个赞吧

分类:题解

评论 (0)

无需注册,留言经站长审核后公开。

加载评论中…

Leetcode 2602:前缀和+二分查找 · Yinchen 的博客