2448. 使数组相等的最小开销

题目描述

给你两个下标从 0 开始的数组 nums 和 cost ，分别包含 n 个 正 整数。

你可以执行下面操作 任意 次：

• 将 nums 中 任意 元素增加或者减小 1 。

对第 i 个元素执行一次操作的开销是 cost[i] 。

请你返回使 nums 中所有元素 相等 的 最少 总开销。

 

示例 1：

输入：nums = [1,3,5,2], cost = [2,3,1,14]
输出：8
解释：我们可以执行以下操作使所有元素变为 2 ：
- 增加第 0 个元素 1 次，开销为 2 。
- 减小第 1 个元素 1 次，开销为 3 。
- 减小第 2 个元素 3 次，开销为 1 + 1 + 1 = 3 。
总开销为 2 + 3 + 3 = 8 。
这是最小开销。

示例 2：

输入：nums = [2,2,2,2,2], cost = [4,2,8,1,3]
输出：0
解释：数组中所有元素已经全部相等，不需要执行额外的操作。

 

提示：

• n == nums.length == cost.length
• 1 <= n <= 105
• 1 <= nums[i], cost[i] <= 106
• 测试用例确保输出不超过 2⁵³-1。

解法

方法一：前缀和 + 排序 + 枚举

我们记数组 nums 所有元素为 $a_1, a_2, \cdots, a_n$，数组 cost 所有元素为 $b_1, b_2, \cdots, b_n$。我们不妨令 $a_1 \leq a_2 \leq \cdots \leq a_n$，即数组 nums 升序排列。

假设我们将数组 nums 中的元素全部变为 $x$，那么我们需要的总开销为：

$$ \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} $$

其中 $k$ 为 $a_1, a_2, \cdots, a_n$ 中小于等于 $x$ 的元素个数。

我们可以使用前缀和的方法，计算 $\sum_{i=1}^{k} b_i$ 和 $\sum_{i=1}^{k} a_ib_i$，以及 $\sum_{i=k+1}^{n}a_ib_i$ 和 $\sum_{i=k+1}^{n}b_i$。

然后我们枚举 $x$，计算上述四个前缀和，得到上述的总开销，取其中的最小值即可。

时间复杂度 $O(n\times \log n)$。其中 $n$ 为数组 nums 的长度。主要是排序的时间复杂度。

Python3JavaC++GoRust

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

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;
    }
}

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;
    }
};

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
}

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
    }
}

方法二：排序 + 中位数

我们还可以把 $b_i$ 看作是 $a_i$ 的出现次数，那么中位数下标是 $\frac{\sum_{i=1}^{n} b_i}{2}$。把所有数变成中位数，一定是最优的。

时间复杂度 $O(n\times \log n)$。其中 $n$ 为数组 nums 的长度。主要是排序的时间复杂度。

相似题目：

• 296. 最佳的碰头地点
• 462. 最少移动次数使数组元素相等 II

Python3JavaC++Go

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)

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;
    }
}

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;
    }
};

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
}

评论