随笔- 578 文章- 0 评论- 61 阅读- 10万

Minimum Cost to Make Array Equal

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

You can do the following operation any number of times:

Increase or decrease any element of the array 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代码如下：

 1 class Solution {
 2 public:
 3     long long minCost(vector<int>& nums, vector<int>& cost) {
 4         int n = nums.size();
 5         vector<int> p;
 6         for (int i = 0; i < n; i++) {
 7             p.push_back(i);
 8         }
 9         sort(p.begin(), p.end(), [&](int i, int j){ // 对nums从小到大排序
10             return nums[i] < nums[j];
11         });
12         vector<long long> sc(n + 1), s(n + 1);
13         for (int i = 1; i <= n; i++) {  // 预处理前缀和
14             sc[i] = sc[i - 1] + cost[p[i - 1]];
15             s[i] = s[i - 1] + 1ll * nums[p[i - 1]] * cost[p[i - 1]];
16         }
17         long long ret = 9e18;
18         for (int i = 1; i <= n; i++) {  // 枚举所有一元线性方程的端点求最小值
19             ret = min(ret, (sc[i] - (sc[n] - sc[i])) * nums[p[i - 1]] + s[n] - s[i] - s[i]);
20         }
21         return ret;
22     }
23 };

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

　　首先有结论，对于问题 $\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代码如下：

 1 class Solution {
 2 public:
 3     long long minCost(vector<int>& nums, vector<int>& cost) {
 4         int n = nums.size();
 5         vector<int> p;
 6         for (int i = 0; i < n; i++) {
 7             p.push_back(i);
 8         }
 9         sort(p.begin(), p.end(), [&](int i, int j){
10             return nums[i] < nums[j];
11         });
12         long long mid = accumulate(cost.begin(), cost.end(), 1ll) - 1 >> 1;
13         long long s = 0, x;
14         for (int i = 0; i < n; i++) {
15             s += cost[p[i]];
16             if (s > mid) {  // 此时中位数就是nums[p[i]]
17                 x = nums[p[i]];
18                 break;
19             }
20         }
21         long long ret = 0;
22         for (int i = 0; i < n; i++) {
23             ret += 1ll * abs(nums[i] - x) * cost[i];
24         }
25         return ret;
26     }
27 };