Mark Elements on Array by Performing Queries

You are given a 0-indexed array nums of size n consisting of positive integers.

You are also given a 2D array queries of size m where queries[i] = [index<sub>i</sub>, k<sub>i</sub>].

Initially all elements of the array are unmarked.

You need to apply m queries on the array in order, where on the i<sup>th</sup> query you do the following:

• Mark the element at index index<sub>i</sub> if it is not already marked.

• Then mark k<sub>i</sub> unmarked elements in the array with the smallest values. If multiple such elements exist, mark the ones with the smallest indices. And if less than k<sub>i</sub> unmarked elements exist, then mark all of them.

Return an array answer of size m where answer[i] is the sum of unmarked elements in the array after the i<sup>th</sup> query.

public class Solution {

    // Method to find the sum of unmarked elements in an array after performing queries
    public long[] unmarkedSumArray(int[] nums, int[][] queries) {
        int n = nums.length;  // Length of the input array
        boolean[] marked = new boolean[n];  // Array to track marked elements
        long[] answer = new long[queries.length];  // Array to store the result for each query
        long unmarkedSum = 0;  // Variable to store the sum of unmarked elements

        // Initialize the sum of all elements in the array
        for (int num : nums) {
            unmarkedSum += num;
        }

        // Min-heap (PriorityQueue) to store pairs of numbers and their indices
        PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> {
            if (a[0] == b[0]) {
                return a[1] - b[1];  // If values are equal, compare by index
            }
            return a[0] - b[0];  // Otherwise, compare by value
        });

        // Populate the min-heap with all elements and their indices
        for (int i = 0; i < n; i++) {
            minHeap.offer(new int[]{nums[i], i});
        }

        // Process each query
        for (int i = 0; i < queries.length; i++) {
            int index = queries[i][0];  // Index to be marked in this query
            int k = queries[i][1];  // Number of additional smallest elements to mark

            // If the element at the given index is not marked, mark it and adjust the sum
            if (!marked[index]) {
                marked[index] = true;
                unmarkedSum -= nums[index];
            }

            // Mark the next 'k' smallest unmarked elements
            for (int j = 0; j < k; j++) {
                // Poll elements from the heap until an unmarked element is found
                while (!minHeap.isEmpty()) {
                    int[] top = minHeap.poll();
                    int value = top[0];  // Value of the element
                    int idx = top[1];  // Index of the element

                    // If the element is unmarked, mark it and adjust the sum
                    if (!marked[idx]) {
                        marked[idx] = true;
                        unmarkedSum -= value;
                        break;  // Exit loop after marking the element
                    }
                }

                // If the heap becomes empty, stop further marking
                if (minHeap.isEmpty()) {
                    break;
                }
            }

            // Store the current sum of unmarked elements after this query
            answer[i] = unmarkedSum;
        }

        // Return the results of all queries
        return answer;
    }
}