Minimum Cost to Make Array Equal

Minimum Cost to Make Array Equal

You are given two 0-indexed arrays nums and cost consisting each of $n$ positive integers.

You can do the following operation any number of times:

• Increase or decrease any element of the array nums by $1$.

The cost of doing one operation on the `ith` element is cost[i] .

Return the minimum total cost such that all the elements of the array nums become equal.

Example 1:

Input: nums = [1,3,5,2], cost = [2,3,1,14]
Output: 8
Explanation: We can make all the elements equal to 2 in the following way:
- Increase the 0th element one time. The cost is 2.
- Decrease the 1st element one time. The cost is 3.
- Decrease the 2nd element three times. The cost is 1 + 1 + 1 = 3.
The total cost is 2 + 3 + 3 = 8.
It can be shown that we cannot make the array equal with a smaller cost.

Example 2:

Input: nums = [2,2,2,2,2], cost = [4,2,8,1,3]
Output: 0
Explanation: All the elements are already equal, so no operations are needed.

Constraints:

• $ n \ \mathrm{==} \ \text{nums.length} \ \mathrm{==} \ \text{cost.length}$
• $1 \leq n \leq {10}^{5}$
• $1 \leq \text{nums}[i], \text{cost}[i] \leq {10}^{6}$

 

解题思路

　　数学题，先把式子写出来，假设最终所有的数字都变成$x$，那么答案就是$$\sum\limits_{i=1}^{n}{|a_i - x|} \cdot c_i$$

　　现在要求这个式子的最小值。假设$a_i$已经从小到大排好序，且$x \in \left[ a_k, a_{k+1} \right)$，那么

\begin{align*}
& \ \ \ \ \ \sum\limits_{i=1}^{n}{|a_i - x|} \cdot c_i \\
&= \sum\limits_{i=1}^{k}{(x - a_i)} \cdot c_i + \sum\limits_{i=k+1}^{n}{(a_i - x)} \cdot c_i \\
&= \sum\limits_{i=1}^{k}{c_i} \cdot x - \sum\limits_{i=1}^{k}{a_i c_i} - \sum\limits_{i=k+1}^{n}{c_i} \cdot x + \sum\limits_{i=k+1}^{n}{a_i c_i} \\
&= \left( {\sum\limits_{i=1}^{k}{c_i} - \sum\limits_{i=k+1}^{n}{c_i}} \right) \cdot x + \sum\limits_{i=k+1}^{n}{a_i c_i} - \sum\limits_{i=1}^{k}{a_i c_i}
\end{align*}

　　记${sc}_{i} = \sum\limits_{j=1}^{i}{c_j}$，$s_{i} = \sum\limits_{j=1}^{i}{a_j c_j}$，那么上式就可以表示成$$\left( 2 \cdot {sc}_{k} - {sc}_{n} \right) \cdot x + \left( s_n - 2 \cdot s_k \right)$$

　　容易发现这个式子是一个一元线性方程，要取得式子的最小值那么$x$一定是取定义域的端点处，即$x = a_k$或$x=a_{k+1}$。因此可以先预处理出来前缀和$sc$和$s$，然后枚举所有的$a_k$，把相应的值带到式子中求最小值。

　　AC代码如下：

class Solution {
public:
    long long minCost(vector<int>& nums, vector<int>& cost) {
        int n = nums.size();
        vector<int> p;
        for (int i = 0; i < n; i++) {
            p.push_back(i);
        }
        sort(p.begin(), p.end(), [&](int i, int j){ // 对nums从小到大排序
            return nums[i] < nums[j];
        });
        vector<long long> sc(n + 1), s(n + 1);
        for (int i = 1; i <= n; i++) {  // 预处理前缀和
            sc[i] = sc[i - 1] + cost[p[i - 1]];
            s[i] = s[i - 1] + 1ll * nums[p[i - 1]] * cost[p[i - 1]];
        }
        long long ret = 9e18;
        for (int i = 1; i <= n; i++) {  // 枚举所有一元线性方程的端点求最小值
            ret = min(ret, (sc[i] - (sc[n] - sc[i])) * nums[p[i - 1]] + s[n] - s[i] - s[i]);
        }
        return ret;
    }
};

　　再给出另外一种做法，是关于中位数的结论，这个还是很难想到的。

　　首先有结论，对于问题$\min \sum\limits_{i=1}^{n}{|a_i-x|}$，当$x = a_ {\left\lfloor \frac{n + 1 \ }{2} \right\rfloor}$，即$x$取$a$的中位数时，式子能够取到最小值（$a$已从小到大排序）。

　　而现在的问题是$\min \sum\limits_{i=1}^{n}{|a_i-x| \cdot c_i}$，$c_i$就相当于带了权值（上面式子的$c_i$均为$1$），但一样可以用上面的方法来求，只需要将$|a_i-x| \cdot c_i$看作是$c_i$个系数为$1$的$|a_i-x|$累加。

　　因此思路就变成了求$\sum\limits_{i=1}^{n}{c_i}$的中位数$\text{mid}$，然后取$x = a_{\text{mid}}$带到式子里算答案就可以了。

　　AC代码如下：

class Solution {
public:
    long long minCost(vector<int>& nums, vector<int>& cost) {
        int n = nums.size();
        vector<int> p;
        for (int i = 0; i < n; i++) {
            p.push_back(i);
        }
        sort(p.begin(), p.end(), [&](int i, int j){
            return nums[i] < nums[j];
        });
        long long mid = accumulate(cost.begin(), cost.end(), 1ll) - 1 >> 1;
        long long s = 0, x;
        for (int i = 0; i < n; i++) {
            s += cost[p[i]];
            if (s > mid) {  // 此时中位数就是nums[p[i]]
                x = nums[p[i]];
                break;
            }
        }
        long long ret = 0;
        for (int i = 0; i < n; i++) {
            ret += 1ll * abs(nums[i] - x) * cost[i];
        }
        return ret;
    }
};

 

参考资料

　　数学 & 枚举：https://leetcode.cn/problems/minimum-cost-to-make-array-equal/solution/by-tsreaper-8hxm/

　　【小羊肖恩】第 316 场力扣周赛题解：https://leetcode.cn/circle/discuss/uO4WuN/view/CUy95z/