LeetCode 1029. Two City Scheduling 【贪心,堆,排序,快速选择】中等

公司计划面试 2n 人。给你一个数组 costs ，其中 costs[i] = [aCosti, bCosti] 。第 i 人飞往 a 市的费用为 aCosti ，飞往 b 市的费用为 bCosti 。

返回将每个人都飞到 a 、b 中某座城市的最低费用，要求每个城市都有 n 人抵达。

示例 1：

输入：costs = [[10,20],[30,200],[400,50],[30,20]]
输出：110
解释：
第一个人去 a 市，费用为 10。
第二个人去 a 市，费用为 30。
第三个人去 b 市，费用为 50。
第四个人去 b 市，费用为 20。
最低总费用为 10 + 30 + 50 + 20 = 110，每个城市都有一半的人在面试。

示例 2：

输入：costs = [[259,770],[448,54],[926,667],[184,139],[840,118],[577,469]]
输出：1859

示例 3：

输入：costs = [[515,563],[451,713],[537,709],[343,819],[855,779],[457,60],[650,359],[631,42]]
输出：3086

提示：

2 * n == costs.length
2 <= costs.length <= 100
costs.length 为偶数
1 <= aCosti, bCosti <= 1000

假设公司首先将这 2 N 2N 2N 个人全部安排飞往 B B B 市，再选出 N N N 个人改变他们的行程，让他们飞往 A A A 市。如果选择改变一个人的行程，那么公司将会额外付出 priceA - priceB 的费用，这个费用可正可负。
[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-d56OGU8I-1686295644365)(https://image-1307616428.cos.ap-beijing.myqcloud.com/Obsidian/202306081038751.png)]
因此最优的方案是，选出 priceA - priceB 最小的 N N N 个人去 A A A 市，其余人去 B B B 市。

解法1 贪心+排序

具体实现来说，这个解法先按照 priceA - priceB 从小到大排序，将前 N N N 个人飞往 A A A 市，其余人飞往 B B B 市，并计算出总费用。

class Solution {
    public int twoCitySchedCost(int[][] costs) {
        // sort by a gain which company has by 
        // sending a person to city A and not to city B
        Arrays.sort(costs, (a, b) -> a[0] - a[1] - (b[0] - b[1]));
        int ans = 0, n = costs.length / 2;
        for (int i = 0; i < n; ++i) ans += costs[i][0] + costs[i + n][1];
        return ans;
    }
}

代码实现简单，分析一下时空复杂度：

时间复杂度： O ( 2 N × log ⁡ 2 N ) O(2N \times \log 2N) O(2N×log2N)
空间复杂度： O ( 1 ) O(1) O(1)

解法2 贪心+堆

对上述解法使用堆进行优化，我们可以使用小顶堆，将差价 priceA - priceB 小的排在堆前面，priceA - priceB 最小的在堆顶，然后从堆中取出 N N N 个当前差价 priceA - priceB 最小的项，在 priceB 的和之上，加上改变行程付出的代价。这种做法要将所有项入堆、并出堆 N N N 个最小的项，因此时间复杂度为 O ( 2 N × log ⁡ 2 N + N log ⁡ 2 N ) O(2N \times \log 2N + N\log 2N) O(2N×log2N+Nlog2N) 。

我们还可使用大顶堆，保持大顶堆的大小为 N N N ：

当大顶堆大小小于 N N N 时直接入堆；
否则看当前人员的差价 priceA - priceB 是否小于大顶堆的堆顶元素，小则先弹出堆顶元素，将当前差价入堆。
这样最后得到的堆就是 N N N 个具有最小 priceA - priceB 的项。
最后，取出堆中所有项，在 priceB 的和之上，加上改变行程付出的代价。

这种做法不用将所有项入堆和出堆，因此降低了时间复杂度。代码实现如下：

class Solution {
    public int twoCitySchedCost(int[][] costs) {
        PriorityQueue<Integer> pq = new PriorityQueue<>((a, b) -> b - a);
        int n = costs.length, halfN = n >> 1;
        int ans = 0;
        for (int i = 0; i < n; ++i) { 
            ans += costs[i][1]; // 所有人都去b城市的费用总和
            pq.offer(costs[i][0] - costs[i][1]);
            if (pq.size() > halfN) pq.poll();
        }
        while (!pq.isEmpty()) ans += pq.poll(); // +diff
        return ans;
    }
}

复杂度分析：

时间复杂度： O ( 2 N × log ⁡ N + N log ⁡ N ) O(2N \times \log N + N\log N) O(2N×logN+NlogN)
空间复杂度： O ( N ) O(N) O(N)

上述写法相当简单易懂，但还可稍微优化：

我们令 a n s ans ans 每次加上 c o s t s [ i ] [ 0 ] costs[i][0] costs[i][0] ，这样 a n s ans ans 就是所有人都去 A A A 市的费用总和；
当堆的大小超过 N N N 时，弹出的都是不去 A A A 市的，用 a n s − p q . p o l l ( ) ans - pq.poll() ans−pq.poll() 即为 a n s − c o s t s [ i ] [ 0 ] + c o s t s [ i ] [ 1 ] ans - costs[i][0] + costs[i][1] ans−costs[i][0]+costs[i][1] ，这样就让他去 B B B 市了。

class Solution {
    public int twoCitySchedCost(int[][] costs) {
        PriorityQueue<Integer> pq = new PriorityQueue<>((a, b) -> b - a);
        int n = costs.length, halfN = n >> 1;
        int ans = 0;
        for (int i = 0; i < n; ++i) { 
            ans += costs[i][0]; // 所有人都去A市的费用总和
            pq.offer(costs[i][0] - costs[i][1]);
            if (pq.size() > halfN) ans -= pq.poll(); // 弹出的都是不去A市的
        } 
        return ans;
    }
}

复杂度分析：

时间复杂度： O ( 2 N × log ⁡ N ) O(2N \times \log N) O(2N×logN)
空间复杂度： O ( N ) O(N) O(N)

解法3 贪心+快速选择

事实上，只需要优化排序部分，就可达到 O ( N ) O(N) O(N) ，严格来说是 O ( 2 N ) O(2N) O(2N) 。优化思路如下：

只需要知道哪一半人去 A A A 更划算；
无需对所有人的差价进行排序，甚至不需要用堆获取「 N N N 个人的最小差价」。
实现：分治思想，利用快速排序的 partition 算法。
- 标准快速排序是全排序，划分后左右都需要排序，则 T ( n ) = 2 T ( n / 2 ) + O ( n ) T(n) = 2T(n/2) +O(n) T(n)=2T(n/2)+O(n) ；
- 本题只需要排一半，则 T ( n ) = T ( n / 2 ) + O ( n ) → O ( 2 n ) T(n) = T(n/2) + O(n) \to O(2n) T(n)=T(n/2)+O(n)→O(2n) 。

class Solution {
    public int twoCitySchedCost(int[][] costs) {
        int ans = 0, n = costs.length;
        int[] arr = new int[n];
        for (int i = 0; i < n; ++i) {
            ans += costs[i][0]; // 都去A市的费用
            arr[i] = costs[i][0] - costs[i][1];
        }
        quickSort(arr, 0, n - 1);
        for (int i = n / 2; i < n; ++i) ans -= arr[i]; // 差价大的,说明去B市更好
        return ans;
    }
    private void quickSort(int[] arr, int l, int r) {
        if (l >= r) return;
        int p = partition(arr, l, r), mid = arr.length / 2;
        if (p == mid) return; // 已经是中点
        if (p > mid) quickSort(arr, l, p - 1);
        else quickSort(arr, p + 1, r);
    } 
    private int partition(int[] arr, int l, int r) {
        int c = l;
        for (int i = l; i < r; ++i)
            if (arr[i] < arr[r]) swap(arr, i, c++);
        swap(arr, r, c);
        return c;
    }
    private void swap(int[] arr, int i, int j) {
        if (i == j) return;
        int swap = arr[i];
        arr[i] = arr[j];
        arr[j] = swap;
    }
}

复杂度分析：

时间复杂度： O ( N ) O(N) O(N)
空间复杂度： O ( N ) O(N) O(N)

memcpy0