6357.使数组元素全部相等的最少操作次数-338

发布时间 2023-03-26 16:42:39作者: huangxk23

使数组元素全部相等的最小操作次数

给你一个正整数数组 nums 。

同时给你一个长度为 m 的整数数组 queries 。第 i 个查询中,你需要将 nums 中所有元素变成 queries[i] 。你可以执行以下操作 任意 次:

将数组里一个元素 增大 或者 减小 1 。
请你返回一个长度为 m 的数组 answer ,其中 answer[i]是将 nums 中所有元素变成 queries[i] 的 最少 操作次数。

注意,每次查询后,数组变回最开始的值。

示例 1:

输入:nums = [3,1,6,8], queries = [1,5]
输出:[14,10]
解释:第一个查询,我们可以执行以下操作:

  • 将 nums[0] 减小 2 次,nums = [1,1,6,8] 。
  • 将 nums[2] 减小 5 次,nums = [1,1,1,8] 。
  • 将 nums[3] 减小 7 次,nums = [1,1,1,1] 。
    第一个查询的总操作次数为 2 + 5 + 7 = 14 。
    第二个查询,我们可以执行以下操作:
  • 将 nums[0] 增大 2 次,nums = [5,1,6,8] 。
  • 将 nums[1] 增大 4 次,nums = [5,5,6,8] 。
  • 将 nums[2] 减小 1 次,nums = [5,5,5,8] 。
  • 将 nums[3] 减小 3 次,nums = [5,5,5,5] 。
    第二个查询的总操作次数为 2 + 4 + 1 + 3 = 10 。
    示例 2:

输入:nums = [2,9,6,3], queries = [10]
输出:[20]
解释:我们可以将数组中所有元素都增大到 10 ,总操作次数为 8 + 1 + 4 + 7 = 20 。

提示:

n == nums.length
m == queries.length
1 <= n, m <= \(10^5\)
1 <= nums[i], queries[i] <= \(10^9\)

来源:力扣(LeetCode)
链接:https://leetcode.cn/problems/minimum-operations-to-make-all-array-elements-equal
著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。

解题思路

朴素的写法实际上是较为容易想到的,就是遍历queries[i]和nums中每一个元素之间的差值,时间复杂度是O(m * n)的,是过不了的。当时写的时候有两个想法,一个是优化queries[i]和nums差值的计算,使得不是O(m)的,另一个是queries[i]和queries[i-1]之间的关系,想办法利用之前查询中的历史信息,使得不用再遍历nums。后来两个都没有想好,其实在这个过程中想到了二分查找,将其划分为<=queries[i]以及>queries[i]的两个部分,但是考虑到计算差值还是要遍历,实际上解决计算差值并不需要多次遍历即可,也就是使用前缀和。将其划分为<queries[i]和>=queries[i]两个部分,可以使用二分查找,那么差值就是queries[i] * l -sum_front + sum_back - queries[i] * r,l表示<queries[i]的个数,r反之。对于求和可以使用前缀和一次性求出,并不需要每个queries[i]都去遍历。

code

class Solution {
public:
    typedef long long int ll;

    int find(vector<int> & nums,int n)
    {
        int l = 0,r = nums.size() -1;

        while(l < r)
        {
            int mid = (l + r) / 2;
            if(nums[mid] >= n) r = mid;
            else l = mid + 1;
        }

        return l;

    }
    vector<long long> minOperations(vector<int>& nums, vector<int>& queries) {
        
        int n = nums.size();

        sort(nums.begin(),nums.end());
        
        vector<long long int> prefix_sum(n,0);

        prefix_sum[0] = nums[0];
        for(int i = 1;i < n;i ++) prefix_sum[i] = nums[i] + prefix_sum[i-1];

        vector<long long int> ans;
        long long int total;
        for(int i = 0;i < queries.size();i ++)
        {
            total = 0;
            int l = find(nums,queries[i]);

            //l >= queries[i]
            //1. l == 0 all nums >= queries[i]
            //2. l == n-1 all nums < queries[i] or nums[l] == queries[i]
            if(l == n-1 && nums[l] <= queries[i])
            { 
               total += (ll)queries[i] * n - prefix_sum[n-1];
            }
            else if(l == 0 && nums[l] >= queries[i])
            {
                
                total += prefix_sum[n-1] - (ll)queries[i] * n;
            }
            else
            {
                total += -prefix_sum[l-1] + (ll)queries[i] * l;
               
                total += -(ll)queries[i] * (n - l) + (prefix_sum[n-1] - prefix_sum[l-1]);
                
            }

            ans.push_back(total);

        }

        return ans;
    }
};