Description#
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 == nums.length == cost.length1 <= n <= 1051 <= nums[i], cost[i] <= 106- Test cases are generated in a way that the output doesn't exceed 253-1
Solutions#
Solution 1: Prefix Sum + Sorting + Enumeration#
Let’s denote the elements of the array nums as $a_1, a_2, \cdots, a_n$ and the elements of the array cost as $b_1, b_2, \cdots, b_n$. We can assume that $a_1 \leq a_2 \leq \cdots \leq a_n$, i.e., the array nums is sorted in ascending order.
Suppose we change all elements in the array nums to $x$, then the total cost we need is:
$$
\begin{aligned}
\sum_{i=1}^{n} \left | a_i-x \right | b_i &= \sum_{i=1}^{k} (x-a_i)b_i + \sum_{i=k+1}^{n} (a_i-x)b_i \
&= x\sum_{i=1}^{k} b_i - \sum_{i=1}^{k} a_ib_i + \sum_{i=k+1}^{n}a_ib_i - x\sum_{i=k+1}^{n}b_i
\end{aligned}
$$
where $k$ is the number of elements in $a_1, a_2, \cdots, a_n$ that are less than or equal to $x$.
We can use the prefix sum method to calculate $\sum_{i=1}^{k} b_i$ and $\sum_{i=1}^{k} a_ib_i$, as well as $\sum_{i=k+1}^{n}a_ib_i$ and $\sum_{i=k+1}^{n}b_i$.
Then we enumerate $x$, calculate the above four prefix sums, get the total cost mentioned above, and take the minimum value.
The time complexity is $O(n\times \log n)$, where $n$ is the length of the array nums. The main time complexity comes from sorting.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
| class Solution:
def minCost(self, nums: List[int], cost: List[int]) -> int:
arr = sorted(zip(nums, cost))
n = len(arr)
f = [0] * (n + 1)
g = [0] * (n + 1)
for i in range(1, n + 1):
a, b = arr[i - 1]
f[i] = f[i - 1] + a * b
g[i] = g[i - 1] + b
ans = inf
for i in range(1, n + 1):
a = arr[i - 1][0]
l = a * g[i - 1] - f[i - 1]
r = f[n] - f[i] - a * (g[n] - g[i])
ans = min(ans, l + r)
return ans
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
| class Solution {
public long minCost(int[] nums, int[] cost) {
int n = nums.length;
int[][] arr = new int[n][2];
for (int i = 0; i < n; ++i) {
arr[i] = new int[] {nums[i], cost[i]};
}
Arrays.sort(arr, (a, b) -> a[0] - b[0]);
long[] f = new long[n + 1];
long[] g = new long[n + 1];
for (int i = 1; i <= n; ++i) {
long a = arr[i - 1][0], b = arr[i - 1][1];
f[i] = f[i - 1] + a * b;
g[i] = g[i - 1] + b;
}
long ans = Long.MAX_VALUE;
for (int i = 1; i <= n; ++i) {
long a = arr[i - 1][0];
long l = a * g[i - 1] - f[i - 1];
long r = f[n] - f[i] - a * (g[n] - g[i]);
ans = Math.min(ans, l + r);
}
return ans;
}
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
| using ll = long long;
class Solution {
public:
long long minCost(vector<int>& nums, vector<int>& cost) {
int n = nums.size();
vector<pair<int, int>> arr(n);
for (int i = 0; i < n; ++i) arr[i] = {nums[i], cost[i]};
sort(arr.begin(), arr.end());
vector<ll> f(n + 1), g(n + 1);
for (int i = 1; i <= n; ++i) {
auto [a, b] = arr[i - 1];
f[i] = f[i - 1] + 1ll * a * b;
g[i] = g[i - 1] + b;
}
ll ans = 1e18;
for (int i = 1; i <= n; ++i) {
auto [a, _] = arr[i - 1];
ll l = 1ll * a * g[i - 1] - f[i - 1];
ll r = f[n] - f[i] - 1ll * a * (g[n] - g[i]);
ans = min(ans, l + r);
}
return ans;
}
};
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
| func minCost(nums []int, cost []int) int64 {
n := len(nums)
type pair struct{ a, b int }
arr := make([]pair, n)
for i, a := range nums {
b := cost[i]
arr[i] = pair{a, b}
}
sort.Slice(arr, func(i, j int) bool { return arr[i].a < arr[j].a })
f := make([]int, n+1)
g := make([]int, n+1)
for i := 1; i <= n; i++ {
a, b := arr[i-1].a, arr[i-1].b
f[i] = f[i-1] + a*b
g[i] = g[i-1] + b
}
var ans int64 = 1e18
for i := 1; i <= n; i++ {
a := arr[i-1].a
l := a*g[i-1] - f[i-1]
r := f[n] - f[i] - a*(g[n]-g[i])
ans = min(ans, int64(l+r))
}
return ans
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
| impl Solution {
#[allow(dead_code)]
pub fn min_cost(nums: Vec<i32>, cost: Vec<i32>) -> i64 {
let mut zip_vec: Vec<_> = nums.into_iter().zip(cost.into_iter()).collect();
// Sort the zip vector based on nums
zip_vec.sort_by(|lhs, rhs| { lhs.0.cmp(&rhs.0) });
let (nums, cost): (Vec<i32>, Vec<i32>) = zip_vec.into_iter().unzip();
let mut sum: i64 = 0;
for &c in &cost {
sum += c as i64;
}
let middle_cost = (sum + 1) / 2;
let mut cur_sum: i64 = 0;
let mut i = 0;
let n = nums.len();
while i < n {
if (cost[i] as i64) + cur_sum >= middle_cost {
break;
}
cur_sum += cost[i] as i64;
i += 1;
}
Self::compute_manhattan_dis(&nums, &cost, nums[i])
}
#[allow(dead_code)]
fn compute_manhattan_dis(v: &Vec<i32>, c: &Vec<i32>, e: i32) -> i64 {
let mut ret = 0;
let n = v.len();
for i in 0..n {
if v[i] == e {
continue;
}
ret += ((v[i] - e).abs() as i64) * (c[i] as i64);
}
ret
}
}
|
We can also consider $b_i$ as the occurrence times of $a_i$, then the index of the median is $\frac{\sum_{i=1}^{n} b_i}{2}$. Changing all numbers to the median is definitely optimal.
The time complexity is $O(n\times \log n)$, where $n$ is the length of the array nums. The main time complexity comes from sorting.
Similar problems:
1
2
3
4
5
6
7
8
9
| class Solution:
def minCost(self, nums: List[int], cost: List[int]) -> int:
arr = sorted(zip(nums, cost))
mid = sum(cost) // 2
s = 0
for x, c in arr:
s += c
if s > mid:
return sum(abs(v - x) * c for v, c in arr)
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
| class Solution {
public long minCost(int[] nums, int[] cost) {
int n = nums.length;
int[][] arr = new int[n][2];
for (int i = 0; i < n; ++i) {
arr[i] = new int[] {nums[i], cost[i]};
}
Arrays.sort(arr, (a, b) -> a[0] - b[0]);
long mid = sum(cost) / 2;
long s = 0, ans = 0;
for (var e : arr) {
int x = e[0], c = e[1];
s += c;
if (s > mid) {
for (var t : arr) {
ans += (long) Math.abs(t[0] - x) * t[1];
}
break;
}
}
return ans;
}
private long sum(int[] arr) {
long s = 0;
for (int v : arr) {
s += v;
}
return s;
}
}
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
| using ll = long long;
class Solution {
public:
long long minCost(vector<int>& nums, vector<int>& cost) {
int n = nums.size();
vector<pair<int, int>> arr(n);
for (int i = 0; i < n; ++i) arr[i] = {nums[i], cost[i]};
sort(arr.begin(), arr.end());
ll mid = accumulate(cost.begin(), cost.end(), 0ll) / 2;
ll s = 0, ans = 0;
for (auto [x, c] : arr) {
s += c;
if (s > mid) {
for (auto [v, d] : arr) {
ans += 1ll * abs(v - x) * d;
}
break;
}
}
return ans;
}
};
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
| func minCost(nums []int, cost []int) int64 {
n := len(nums)
type pair struct{ a, b int }
arr := make([]pair, n)
mid := 0
for i, a := range nums {
b := cost[i]
mid += b
arr[i] = pair{a, b}
}
mid /= 2
sort.Slice(arr, func(i, j int) bool { return arr[i].a < arr[j].a })
s, ans := 0, 0
for _, e := range arr {
x, c := e.a, e.b
s += c
if s > mid {
for _, t := range arr {
ans += abs(t.a-x) * t.b
}
break
}
}
return int64(ans)
}
func abs(x int) int {
if x < 0 {
return -x
}
return x
}
|
