[LeetCode] 1029 | Two City Scheduling | Python

---

문제 링크: https://leetcode.com/problems/two-city-scheduling/

---

문제

A company is planning to interview 2n people. Given the array costs where costs[i] = [aCosti, bCosti], the cost of flying the ith person to city a is aCosti, and the cost of flying the ith person to city b is bCosti.

Return the minimum cost to fly every person to a city such that exactly n people arrive in each city.

---

조건

• 2 * n == costs.length
• 2 <= costs.length <= 100
• costs.length is even.
• 1 <= aCosti, bCosti <= 1000

---

내 풀이

costs안의 원소를 [a, b]로 표현한다면, a<b인 것들은 모두 A로, 나머지는 모두 B로 보내놓는다.

A와 B 중 더 많은 원소를 포함한 것을 더 작은 것으로 보내줘야한다. (ex. len(A)=4, len(B)=2일 때, A에서 B로 하나 보내야 함)

어떤 원소를 보내는지는 [a, b]일 때, a와 b의 차가 제일 적은 것부터 보낸다.

그래야 전체 cost가 크게 증가하지 않기 때문이다.

---

코드

class Solution:
    def twoCitySchedCost(self, costs: List[List[int]]) -> int:
        A, B = [], []
        for cost in costs:
            a, b = cost
            if a < b:
                A.append(cost)
            else:
                B.append(cost)
                
        A = sorted(A, key=lambda a: a[1]-a[0])
        B = sorted(B, key=lambda a: a[0]-a[1])
        
        nA, nB, nC = len(A), len(B), int(len(costs)/2)
        if nA > nB:
            for n in range(nC-nB):
                B.append(A[n])
            A = A[nC-nB:]
        else:
            for n in range(nC-nA):
                A.append(B[n])
            B = B[nC-nA:]
        
        ans = 0
        for n in range(len(A)):
            ans += A[n][0]
        for n in range(len(B)):
            ans += B[n][1]
        
        return ans