[牛客题霸-高频算法面试题]一样的水

题目描述

有n个水桶，第i个水桶里面水的体积为Ai，你可以用1秒时间向一个桶里添加1体积的水。

有q次询问，每次询问一个整数pi,你需要求出使其中pi个桶中水的体积相同所花费的最少时间。

对于一次询问如果有多种方案，则采用使最终pi个桶中水的体积最小的方案。

示例1

输入：4,3,[1,2,3,4],[2,2,4]

输出：[1,0,5]

说明：

第一次：花费一秒变为 2 2 3 4

第二次：已经存在两个水的体积一样的桶

第三次：花费五秒从2 2 3 4变为4 4 4 4

 

思路

先来看一下题干，让我们求使最终pi个桶中水的体积最小的方案，也就是说每次向水较少的桶中注水，因此考虑先给n个水桶按含水体积从小到大排个序。要求使其中pi个桶中水的体积相同，对这句话的理解是在排好序的数组中连续pi个数相同。解题的总思路是：每次询问时枚举所有可能的情况，找到花费最小的位置存入res[]中，确定最佳方案后再向桶里注水（修改桶中水的值，即修改数组值）。

 

下面结合例子具体来说一说一次询问时的过程：

• 一共有6个水桶，水桶中初始水的体积为[9, 2, 6, 1, 8, 3]，排序后的数组为[1, 2, 3, 6, 8, 9]，假设某次询问时pi=4，min表示每次询问时最少花费的时间，min_begin表示使花费时间最少的方案的起始水桶的位置
• 可以想象有一个固定大小为4的滑动窗口，从begin=0位置开始向右滑动，sum表示pi个桶中原始水的体积，total表示注水后pi个桶中水的体积
• 第一次：sum=a[0]+a[1]+a[2]+a[3]=12，注水后前4个桶中的水的体积都变成了a[3]=6，所以total=pi*6=24；此时min=total-sum=12，min_begin=0
• 第二次：sum=a[1]+a[2]+a[3]+a[4]=19，注水后第2~5个桶中的水的体积都变成了a[4]=8，所以total=pi*8=32；此时total-sum=13>12，不需要更新min和min_begin
• 第三次：sum=a[2]+a[3]+a[4]+a[5]=26，注水后第3~6个桶中的水的体积都变成了a[5]=9，所以total=pi*9=36；此时total-sum=10<12，更新min=10，min_begin=2
• 本次询问结束。当pi=4时最小的方案是第三次中的方案，即使a[2]到a[5]桶中的水都为9，将min=10存入res[i]中，然后遍历下标为min_begin到min_begin+pi-1的区间并将该区间内的值修改为a[min_begin+pi-1]=9

Tips：

• 计算sum时可以利用滑动窗口的特点，去掉移出窗口的值然后加上新移进窗口的值，即sum=sum-a[begin-1]+a[begin+pi-1]，其中begin为当前方案中第一个桶的位置
• 询问数组p中可能会有重复的值，或是不需要再注水的值（当前n个水桶中含有相同体积水的桶的数目大于pi），这种情况下直接跳过。可以用一个变量equal来记录n个桶中含有相同水的桶的个数，每一轮询问结束后更新equal为Math.max(equal, pi)，下一次询问前先判断该次询问的pi和equal的大小，如果pi<=equal，说明无需注水，另res[i]=0然后直接进入下一轮。

JAVA代码

public class Solution {
    /**
     * 
     * @param n int整型 水桶的个数
     * @param q int整型 询问的次数
     * @param a int整型一维数组 n个水桶中初始水的体积
     * @param p int整型一维数组 每次询问的值
     * @return int整型一维数组
     */
    public int[] solve (int n, int q, int[] a, int[] p) {
        int[] res = new int[q];
        Arrays.sort(a);
        int equal = 0;
        for(int i = 0; i < q; i++) {
            if(p[i] <= equal) {
                res[i] = 0;
                continue;
            }
            int sum = 0;
            for(int j = 0; j < p[i]; j++) {
                sum += a[j];
            }
            int min = a[p[i] - 1] * p[i] - sum, min_begin = 0;;
            for(int begin = 1; begin <= n - p[i]; begin++) {
                sum = sum - a[begin - 1] + a[begin + p[i] - 1];
                if(a[begin + p[i] - 1] * p[i] - sum < min) {
                    min = a[begin + p[i] - 1] * p[i] - sum;
                    min_begin = begin;
                }
            }
            res[i] = min;
            equal = Math.max(equal, p[i]);
            for(int j = min_begin; j < min_begin + p[i]; j++) {
                a[j] = a[min_begin + p[i] - 1];
            }
        }
        return res;
    }
}