[Leetcode] 2611. Mice and Cheese (Medium)

概述

題目

https://leetcode.com/problems/mice-and-cheese/
給出兩個長度相等數字 array，要求拿取一邊的數字就不能拿另一邊的對位數字，共計在第一組數組中取 k 次之最大值

心得

從一邊拿、一邊不拿然後求最大，可判斷盡量拿兩邊差值最大到漸小可最大化結果，操作基本 Greedy，但可分 Sort 或 Heap 來解

Sort

思路

先確認差值，再由大到小排序取第一數組中的 k 個，剩下 n-k 個由第二數組取得

程式

class Solution(object):
    def miceAndCheese(self, reward1, reward2, k):
        """
        :type reward1: List[int]
        :type reward2: List[int]
        :type k: int
        :rtype: int
        """
        r1, r2 = reward1, reward2
        n = len(r1)
        r3 = [[r1[i] - r2[i], i] for i in range(n)]

        # 兩組數組差值大到小排序
        r4 = sorted(r3, reverse=1)
        cnt = 0
        i = 0
        for dif, idx in r4:
            if i < k:
                cnt += r1[idx]
            else:
                cnt += r2[idx]
            i += 1
        return cnt

Complexity

Time Complexity: O(n log n)
有排序故 n log n

Space Complexity: O(1)
基本上無使用到額外空間

Heap

思路

基本同 Sort，只是取法為 Heap 操作之

程式

from heapq import heappush as hh
from heapq import heappop as hp
class Solution(object):
    def miceAndCheese(self, reward1, reward2, k):
        """
        :type reward1: List[int]
        :type reward2: List[int]
        :type k: int
        :rtype: int
        """
        r1, r2 = reward1, reward2
        
        # 放入 Heap 後取出
        h = []
        n = len(r1)
        for i in range(n):
            hh(h, [r2[i]-r1[i], r1[i], r2[i]])
        
        cnt = 0
        for i in range(k):
            dif, c1, c2 = hp(h)
            cnt += c1
        
        for i in range(n-k):
            dif, c1, c2 = hp(h)
            cnt += c2

        return cnt

Complexity

Time Complexity: O(n log n)
有 Heap 故 n log n

Space Complexity: O(1)
基本上無使用到額外空間

Lee215

思路

Lee215 大神一行解

程式

class Solution(object):
    def miceAndCheese(self, reward1, reward2, k):
        """
        :type reward1: List[int]
        :type reward2: List[int]
        :type k: int
        :rtype: int
        """
        return sum(reward2) + sum(nlargest(k, (a-b for a, b in zip(reward1, reward2))))