Leetcode 1029 两地调度优化解法（附OR-Tools和PuLP代码）

Leetcode 1029. 两地调度 (medium)

公司计划面试 2N 人。第 i 人飞往 A 市的费用为 costs[i][0]，飞往 B 市的费用为 costs[i][1]。

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

示例：

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

提示：

1 <= costs.length <= 100 costs.length 为偶数 1 <= costs[i][0], costs[i][1] <= 1000

链接：https://leetcode-cn.com/problems/two-city-scheduling

暴力枚举法

最直接的方式是暴力枚举出所有分组的可能。因为 2N 个人平均分成两组，总数为 \({2n \choose n}\)，是 n 的指数级数量。在文章24 点游戏算法题的 Python 函数式实现: 学用itertools，yield，yield from 巧刷题，我们展示如何调用 Python 的 itertools包，这里，我们也用同样的方式产生 [0, 2N] 的所有集合大小为N的可能（保存在left_set_list中），再遍历找到最小值即可。当然，这种解法会TLE，只是举个例子来体会一下暴力做法。

{linenos

import math
from typing import List

class Solution:
    def twoCitySchedCost(self, costs: List[List[int]]) -> int:
        L = range(len(costs))
        from itertools import combinations
        left_set_list = [set(c) for c in combinations(list(L), len(L)//2)]

        min_total = math.inf
        for left_set in left_set_list:
            cost = 0
            for i in L:
                is_left = 1 if i in left_set else 0
                cost += costs[i][is_left]
            min_total = min(min_total, cost)

        return min_total

O(N) AC解法

对于组合优化问题来说，例如TSP问题（解法链接 TSP问题从DP算法到深度学习1：递归DP方法 AC AIZU TSP问题），一般都是 NP-Hard问题，意味着没有多项次复杂度的解法。但是这个问题比较特殊，它增加了一个特定条件：去城市A和城市B的人数相同，也就是我们已经知道两个分组的数量是一样的。我们仔细思考一下这个意味着什么？考虑只有四个人的小规模情况，如果让你来手动规划，你一定不会穷举出所有两两分组的可能，而是比较人与人相对的两个城市的cost差。举个例子，有如下四个人的costs

0 A:3,  B:1
1 A:99, B:100
2 A:2,  B:2
3 A:3,  B:3

虽然1号人去城市A（99）cost 很大，但是相比较他去B（100）来说，可以省下 100-99 = 1 的钱，这个钱比0号人去B不去A省下的钱 3-1 = 2 还要多，因此你一定会选择让1号人去A而让0号人去B。

有了这个想法，再整理一下，就会发现让某人去哪个城市和他去两个城市的cost 差 $ C_a - C_b$相关，如果这个值越大，那么他越应该去B。但是最后决定他是否去B取决于他的差在所有人中的排名，由于两组人数相等，因此差能大到排在前一半，则他就去B，在后一半就去A。

按照这个思路，很快能写出代码，代码写法有很多，下面略举一例。代码中由于用到排序，复杂度为 \(O(N \cdot \log(N))\) 。这里补充一点，理论上只需找数组中位数的值即可，最好的时间复杂度是 \(O(N)\)。

代码实现上，cost_diff_list 将每个人的在原数组的index 和他的cost差组成 pair。即

[(0, cost_0), (1, cost_1), ... ]

这样我们可以将这个数组按照cost排序，排完序后前面N个元素属于B城市。

{linenos

# AC
# Runtime: 36 ms, faster than 87.77% of Python3 online submissions
# Memory Usage: 14.5 MB, less than 14.84% of Python3 online
from typing import List

class Solution:
    def twoCitySchedCost(self, costs: List[List[int]]) -> int:
        L = range(len(costs))
        cost_diff_lst = [(i, costs[i][0] - costs[i][1]) for i in L]
        cost_diff_lst.sort(key=lambda x: x[1])

        total_cost = 0
        for c, (idx, _) in enumerate(cost_diff_lst):
            is_left = 0 if c < len(L) // 2 else 1
            total_cost += costs[idx][is_left]

        return total_cost

转换成整数规划问题

这个问题对于略有算法经验的人来说，很类似于背包问题。它们都需要回答N个物品取或者不取，并同时最大最小化总cost。区别在它们的约束条件不一样。这道题的约束是去取（去城市A）和不取（去城市B）的数量一样。这一类问题即 integer programming，即整数规划。下面我们选取两个比较流行的优化库来展示如何调包解这道题。

首先我们先来formulate这个问题，因为需要表达两个约束条件，我们将每个人的状态分成是否去A和是否去B两个变量。

x[i-th-person][0]: boolean 表示是否去 city a
x[i-th-person][1]: boolean 表示是否去 city b

这样，问题转换成如下优化模型

\[ \begin{array}{rrclcl} \displaystyle \min_{x} & costs[i][0] \cdot x[i][0] + costs[i][1] \cdot x[i][1]\\ \textrm{s.t.} & x[i][0] + x[i][1] =1\\ &x[i][0] + x[i][1] + ... =N \\ \end{array} \]

Google OR-Tools

Google OR-Tools 是业界最好的优化库，下面为调用代码，由于直接对应于上面的数学优化问题，不做赘述。当然 Leetcode上不支持这些第三方的库，肯定也不能AC。

{linenos

from ortools.sat.python import cp_model

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

I = range(len(costs))

model = cp_model.CpModel()
x = []
total_cost = model.NewIntVar(0, 10000, 'total_cost')
for i in I:
    t = []
    for j in range(2):
        t.append(model.NewBoolVar('x[%i,%i]' % (i, j)))
    x.append(t)

# Constraints
# Each person must be assigned to at exact one city
[model.Add(sum(x[i][j] for j in range(2)) == 1) for i in I]
# equal number of person assigned to two cities
model.Add(sum(x[i][0] for i in I) == (len(I) // 2))

# Total cost
model.Add(total_cost == sum(x[i][0] * costs[i][0] + x[i][1] * costs[i][1] for i in I))
model.Minimize(total_cost)

solver = cp_model.CpSolver()
status = solver.Solve(model)

if status == cp_model.OPTIMAL:
    print('Total min cost = %i' % solver.ObjectiveValue())
    print()
    for i in I:
        for j in range(2):
            if solver.Value(x[i][j]) == 1:
                print('People ', i, ' assigned to city ', j, '  Cost = ', costs[i][j])

完整代码可以从我的github下载。

https://github.com/MyEncyclopedia/leetcode/blob/master/1029_Two_City_Scheduling/1029_ortool.py

PuLP

类似的，另一种流行 python 优化库 PuLP 的代码为

{linenos

import pulp

costs = [[259,770],[448,54],[926,667],[184,139],[840,118],[577,469]]   # 1859


I = range(len(costs))

items=[i for i in I]
city_a = pulp.LpVariable.dicts('left', items, 0, 1, pulp.LpBinary)
city_b = pulp.LpVariable.dicts('right', items, 0, 1, pulp.LpBinary)

m = pulp.LpProblem("Two Cities", pulp.LpMinimize)

m += pulp.lpSum((costs[i][0] * city_a[i] + costs[i][1] * city_b[i]) for i in items)

# Constraints
# Each person must be assigned to at exact one city
for i in I:
    m += pulp.lpSum([city_a[i] + city_b[i]]) == 1
# create a binary variable to state that a table setting is used
m += pulp.lpSum(city_a[i] for i in I) == (len(I) // 2)

m.solve()

total = 0
for i in I:
    if city_a[i].value() == 1.0:
        total += costs[i][0]
    else:
        total += costs[i][1]

print("Total cost {}".format(total))

代码地址为

https://github.com/MyEncyclopedia/leetcode/blob/master/1029_Two_City_Scheduling/1029_pulp.py