LeetCode 3080. Mark Elements on Array by Performing Queries

雙周賽 126。

題目

輸入長度 n 的正整數陣列 nums。

另外輸入長度 m 二維整數陣列 queries，其中 queries[i] = [index_i, k_i]。

最初，nums 中的所有元素都是未標記的。
你必須依序執行 m 次查詢，每次查詢你必須：

• 若 index_i 未標記，則將其標記
• 從未標記的元素中，選擇 k_i 個最小值標記。若存在多個最小值，則選擇索引較小的。若剩餘不足 k_i 個，則全部標記

回傳長度 m 陣列 answer，其中 answer[i] 代表第 i 次查詢後，未標記的元素總和。

解法

看到取最小值，就想到用 min heap 了。
以 [value, index] 為鍵值，將所有元素裝進 heap。同時維護一個 vis 陣列表示元素是否標記過。

題目要求的是未標記的總和，所以要先算出原本的總和，每標記一次就從中扣除。
對於每次查詢，先檢查指定的標記位置，然後取 k 個最小值標記。

時間複雜度 O(N log N)。
空間複雜度 O(N)，答案空間不計入。

class Solution:
    def unmarkedSumArray(self, nums: List[int], queries: List[List[int]]) -> List[int]:
        N = len(nums)
        h = []
        for i, x in enumerate(nums):
            heappush(h, [x, i])
            
        sm = sum(nums)
        vis = [False] * N
        ans = []
        for qi, k in queries:
            if not vis[qi]:
                vis[qi] = True
                sm -= nums[qi]
            
            while h and k:
                v, i = heappop(h)
                if not vis[i]:
                    vis[i] = True
                    k -= 1
                    sm -= v
                    
            ans.append(sm)
            
        return ans

也可以不用 heap。
直接維護一個索引陣列，依照 [value, index] 排序後即可。

class Solution:
    def unmarkedSumArray(self, nums: List[int], queries: List[List[int]]) -> List[int]:
        N = len(nums)
        a = sorted(range(N), key=lambda x:(nums[x], x))
        j = 0
        
        sm = sum(nums)
        vis = [False] * N
        ans = []
        for qi, k in queries:
            if not vis[qi]:
                vis[qi] = True
                sm -= nums[qi]
            
            while j < N and k:
                i = a[j]
                j += 1
                if not vis[i]:
                    vis[i] = True
                    k -= 1
                    sm -= nums[i]
                    
            ans.append(sm)
            
        return ans