We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

Mark Elements on Array by Performing Queries

Number: 3306

Difficulty: Medium

Paid? No

Companies: Barclays


Problem Description

You are given a 0-indexed array nums of size n consisting of positive integers and a 2D array queries where each query is formatted as [index, k]. Initially all elements are unmarked. For the iᵗʰ query, you first mark the element at the specified index (if it was not already marked) and then mark k unmarked elements with the smallest values (if there is a tie, select the one with the smallest index). After processing each query, return the sum of the unmarked elements in the array.


Key Insights

  • Use a min-heap (priority queue) to efficiently access the smallest unmarked element based on value (and index for ties).
  • Maintain a boolean marker array to track if an element has been marked.
  • Keep a running sum of unmarked elements, subtracting the value when an element is marked.
  • Employ lazy deletion in the heap: when extracting the minimum, check if it is already marked, and skip it if so.

Space and Time Complexity

Time Complexity: O(n log n + total_marked * log n) where total_marked is at most n. Space Complexity: O(n) for the heap and the marker array.


Solution

The approach simulates each query using a min-heap data structure that stores tuples of (value, index) for each unmarked element. Initially, the heap contains all elements. For every query:

  1. Mark the queried index (if not already marked) and update the running unmarked sum.
  2. Extract up to k elements from the heap that are unmarked, marking them and updating the sum.
  3. Record the current unmarked sum after processing the query. This simulation strategy, enhanced with lazy deletions in the heap, ensures we always pick the smallest available unmarked element efficiently.

Code Solutions

import heapq

def mark_elements(nums, queries):
    n = len(nums)
    marked = [False] * n  # Track marked indices
    heap = [(nums[i], i) for i in range(n)]  # Heap with (value, index) tuples
    heapq.heapify(heap)
    
    unmarked_sum = sum(nums)  # Running sum of unmarked elements
    answer = []
    
    for idx, k in queries:
        # Mark the queried index if it's not already marked
        if not marked[idx]:
            marked[idx] = True
            unmarked_sum -= nums[idx]
        
        # Mark up to k smallest unmarked elements from the heap
        count = 0
        while count < k and heap:
            value, i = heapq.heappop(heap)
            if not marked[i]:
                marked[i] = True
                unmarked_sum -= value
                count += 1
        
        answer.append(unmarked_sum)
    
    return answer

# Example usage:
nums = [1,2,2,1,2,3,1]
queries = [[1,2],[3,3],[4,2]]
print(mark_elements(nums, queries))  # Output: [8, 3, 0]
← Back to All Questions