Mark Elements on Array by Performing Queries

Med

#2736Time: O(m * n log n). For each of the `m` queries, we scan the `n` elements to find unmarked ones (O(n)) and then sort them (O(n log n)). This makes the total time complexity dominated by the sorting step within the loop.Space: O(n + m). We need O(n) space for the `marked` array and the temporary list of unmarked elements. The answer array requires O(m) space.2 companies

Algorithms

Sorting

Data structures

Array Hash Table Heap (Priority Queue)

Companies

Barclays Samsung

Prompt

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_i, k_i].

Initially all elements of the array are unmarked.

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

Mark the element at index index_i if it is not already marked.
Then mark k_i 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_i 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^th query.

Example 1:

Input: nums = [1,2,2,1,2,3,1], queries = [[1,2],[3,3],[4,2]]

Output: [8,3,0]

Explanation:

We do the following queries on the array:

Mark the element at index 1, and 2 of the smallest unmarked elements with the smallest indices if they exist, the marked elements now are nums = [1,2,2,1,2,3,1]. The sum of unmarked elements is 2 + 2 + 3 + 1 = 8.
Mark the element at index 3, since it is already marked we skip it. Then we mark 3 of the smallest unmarked elements with the smallest indices, the marked elements now are nums = [1,2,2,1,2,3,1]. The sum of unmarked elements is 3.
Mark the element at index 4, since it is already marked we skip it. Then we mark 2 of the smallest unmarked elements with the smallest indices if they exist, the marked elements now are nums = [1,2,2,1,2,3,1]. The sum of unmarked elements is 0.

Example 2:

Input: nums = [1,4,2,3], queries = [[0,1]]

Output: [7]

Explanation: We do one query which is mark the element at index 0 and mark the smallest element among unmarked elements. The marked elements will be nums = [1,4,2,3], and the sum of unmarked elements is 4 + 3 = 7.

Constraints:

n == nums.length
m == queries.length
1 <= m <= n <= 10⁵
1 <= nums[i] <= 10⁵
queries[i].length == 2
0 <= index_i, k_i <= n - 1

Approaches

3 approaches with complexity analysis and trade-offs.

This approach directly simulates the process described in the problem. For each query, it first marks the specified index. Then, to find the k smallest unmarked elements, it gathers all currently unmarked elements into a temporary list, sorts this list based on value and then index, and marks the top k elements from the sorted list. The sum of unmarked elements is updated at each step.

Algorithm

Initialize a boolean array marked of size n to keep track of marked elements, initially all false.
Calculate the initial sum of all elements in nums and store it in totalSum.
For each query [index, k]:
1. Check if nums[index] is already marked. If not, mark it by setting marked[index] = true and subtract nums[index] from totalSum.
2. To find the k smallest unmarked elements, create a temporary list of pairs (value, original_index) for all elements where marked[original_index] is false.
3. Sort this temporary list. The primary sorting key is the value, and the secondary key is the index.
4. Iterate through the first k elements of the sorted list (or fewer if the list is smaller than k).
5. For each of these elements, mark its original index in the marked array and subtract its value from totalSum.
6. After processing both marking steps, add the current totalSum to the answer array.
Return the answer array.

Walkthrough

We maintain a boolean array marked to keep track of marked elements and a variable totalSum for the sum of unmarked elements. Initially, marked is all false, and totalSum is the total sum of nums. For each query [index, k], we first handle marking by index. If marked[index] is false, we set it to true and update totalSum. Then, for marking by value, we build a list of all unmarked elements with their values and original indices. This list is sorted to find the smallest elements. We then iterate through the top k elements of this sorted list, mark them, and update totalSum. Finally, the updated totalSum is recorded as the answer for the current query.

1class Solution {2    public long[] unmarkedSumArray(int[] nums, int[][] queries) {3        int n = nums.length;4        int m = queries.length;5        long[] ans = new long[m];6        boolean[] marked = new boolean[n];7        long totalSum = 0;8        for (int num : nums) {9            totalSum += num;10        }11 12        for (int i = 0; i < m; i++) {13            int index = queries[i][0];14            int k = queries[i][1];15 16            // Mark by index17            if (!marked[index]) {18                marked[index] = true;19                totalSum -= nums[index];20            }21 22            // Mark k smallest unmarked elements23            List<int[]> unmarkedElements = new ArrayList<>();24            for (int j = 0; j < n; j++) {25                if (!marked[j]) {26                    unmarkedElements.add(new int[]{nums[j], j});27                }28            }29 30            unmarkedElements.sort((a, b) -> {31                if (a[0] != b[0]) {32                    return Integer.compare(a[0], b[0]);33                } else {34                    return Integer.compare(a[1], b[1]);35                }36            });37 38            int count = 0;39            for (int[] element : unmarkedElements) {40                if (count >= k) break;41                int val = element[0];42                int idx = element[1];43                if (!marked[idx]) {44                    marked[idx] = true;45                    totalSum -= val;46                    count++;47                }48            }49            ans[i] = totalSum;50        }51 52        return ans;53    }54}

Complexity

Time

O(m * n log n). For each of the `m` queries, we scan the `n` elements to find unmarked ones (O(n)) and then sort them (O(n log n)). This makes the total time complexity dominated by the sorting step within the loop.

Space

O(n + m). We need O(n) space for the `marked` array and the temporary list of unmarked elements. The answer array requires O(m) space.

Trade-offs

Pros

Simple to understand and follows the problem statement directly.

Cons

Extremely inefficient due to repeated scanning and sorting of the array.
Will result in a 'Time Limit Exceeded' error for the given constraints.

This approach significantly optimizes the process of finding the k smallest unmarked elements. Instead of repeatedly sorting a list, we use a min-heap (Priority Queue in Java) to maintain all elements in a sorted manner throughout the process. This allows for efficient retrieval of the smallest unmarked elements.

Algorithm

Initialize a marked boolean array and calculate the totalSum as before.
Create a min-heap (Priority Queue) and insert all elements from nums as pairs of (value, index).
The heap should be ordered by value first, then by index for tie-breaking.
For each query [index, k]:
1. Mark by index: If !marked[index], set marked[index] = true and totalSum -= nums[index].
2. Mark by value: Initialize a counter markedCount = 0.
3. While markedCount < k and the heap is not empty, extract the minimum element (val, idx) from the heap.
4. If marked[idx] is true, it means this element was already marked (e.g., by index in a previous step), so we discard it and continue.
5. If !marked[idx], we mark it (marked[idx] = true), subtract its value from totalSum, and increment markedCount.
6. Store the final totalSum for the query.
Return the answer array.

Walkthrough

We start by populating a min-heap with all elements of nums, storing them as (value, index) pairs. The heap's comparator ensures that elements with smaller values, or smaller indices in case of a tie, have higher priority. We also maintain a marked boolean array and a totalSum of unmarked elements. For each query, we first mark the element at the given index and update the sum. Then, we extract elements from the heap k times. A crucial step is to check if an element extracted from the heap is already marked. If it is, we ignore it and proceed to the next, as it must have been marked by its index in a prior operation. If not marked, we mark it, update the sum, and continue until k new elements are marked or the heap is empty.

1class Solution {2    public long[] unmarkedSumArray(int[] nums, int[][] queries) {3        int n = nums.length;4        int m = queries.length;5        long[] ans = new long[m];6        boolean[] marked = new boolean[n];7        long totalSum = 0;8        for (int num : nums) {9            totalSum += num;10        }11 12        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> {13            if (a[0] != b[0]) {14                return Integer.compare(a[0], b[0]);15            } else {16                return Integer.compare(a[1], b[1]);17            }18        });19 20        for (int i = 0; i < n; i++) {21            pq.offer(new int[]{nums[i], i});22        }23 24        for (int i = 0; i < m; i++) {25            int index = queries[i][0];26            int k = queries[i][1];27 28            // Mark by index29            if (!marked[index]) {30                marked[index] = true;31                totalSum -= nums[index];32            }33 34            // Mark k smallest unmarked elements35            int count = 0;36            while (count < k && !pq.isEmpty()) {37                int[] top = pq.poll();38                int val = top[0];39                int idx = top[1];40 41                if (!marked[idx]) {42                    marked[idx] = true;43                    totalSum -= val;44                    count++;45                }46            }47            ans[i] = totalSum;48        }49 50        return ans;51    }52}

Complexity

Time

O((n+m) log n) or more tightly O(n log n + m). Building the heap takes O(n log n). Each of the `n` elements is extracted from the heap at most once across all `m` queries. The total time for all heap extractions is O(n log n). The main loop runs `m` times.

Space

O(n + m). The min-heap stores `n` elements, the `marked` array takes O(n) space, and the answer array takes O(m) space.

Trade-offs

Pros

Much more efficient than the brute-force approach.
A standard and robust solution for problems requiring repeated extraction of minimum/maximum elements.

Cons

While efficient, the O(n log n) complexity might be suboptimal if the range of values in nums is small, where a linear time solution could exist.

This is the most efficient approach, achieving linear time complexity by leveraging the constraint on the range of values in nums. It avoids general-purpose sorting by grouping elements by their value, similar to a counting sort. This allows finding the smallest unmarked elements by simply iterating through values in increasing order.

Algorithm

Initialize a marked boolean array and totalSum.
Pre-process nums by grouping indices by their value. Use an array of lists, valueToIndices, where valueToIndices[v] contains a sorted list of indices i where nums[i] == v.
Maintain a valuePointer to track the smallest value we are currently considering, and an array of indexPointers to track the next available index for each value's list.
For each query [index, k]:
1. Mark by index as usual, updating marked and totalSum.
2. Mark by value: Start from the current valuePointer. Iterate through the list of indices for this value, starting from its indexPointer.
3. For each index, if it's not marked, mark it, update totalSum, and increment a counter for the k elements.
4. If the list for the current value is exhausted, increment valuePointer to move to the next smallest value.
5. Continue until k elements are marked or all elements are processed.
6. Store the totalSum.
Return the answer array.

Walkthrough

The core idea is to pre-process the array to group indices by value. We can use an array of lists, say valueToIndices, where the index of the array corresponds to a value from nums, and the list at that index stores all original indices where this value appears. By populating this structure by iterating through nums from left to right, the lists of indices will be naturally sorted. We then use pointers to traverse this structure efficiently. A valuePointer keeps track of the smallest value we need to consider, and an indexPointer for each value's list tracks how many of its occurrences we've already marked. During a query, to mark k elements, we use these pointers to find the next available unmarked elements in O(1) amortized time, starting from the smallest values and respecting the smallest index rule.

1class Solution {2    public long[] unmarkedSumArray(int[] nums, int[][] queries) {3        int n = nums.length;4        int m = queries.length;5        long[] ans = new long[m];6        boolean[] marked = new boolean[n];7        long totalSum = 0;8 9        List<Integer>[] valueToIndices = new List[100001];10        for (int i = 0; i < n; i++) {11            totalSum += nums[i];12            if (valueToIndices[nums[i]] == null) {13                valueToIndices[nums[i]] = new ArrayList<>();14            }15            valueToIndices[nums[i]].add(i);16        }17 18        int[] indexPointers = new int[100001];19        int valuePointer = 1;20 21        for (int i = 0; i < m; i++) {22            int index = queries[i][0];23            int k = queries[i][1];24 25            // Mark by index26            if (!marked[index]) {27                marked[index] = true;28                totalSum -= nums[index];29            }30 31            // Mark k smallest unmarked elements32            int count = 0;33            while (count < k && valuePointer < valueToIndices.length) {34                List<Integer> indices = valueToIndices[valuePointer];35                if (indices == null || indexPointers[valuePointer] >= indices.size()) {36                    valuePointer++;37                    continue;38                }39 40                int originalIndex = indices.get(indexPointers[valuePointer]);41                if (!marked[originalIndex]) {42                    marked[originalIndex] = true;43                    totalSum -= valuePointer;44                    count++;45                }46                indexPointers[valuePointer]++;47            }48            ans[i] = totalSum;49        }50 51        return ans;52    }53}

Complexity

Time

O(n + m + W), where `W` is the maximum value in `nums`. O(n) for pre-processing. The query processing takes O(m + n + W) in total because each pointer (`valuePointer` and `indexPointers`) only moves forward, traversing each element and value at most once across all queries.

Space

O(n + m + W), where `W` is the maximum possible value in `nums`. O(n) for the `marked` array and storing all indices, O(W) for `valueToIndices` and `indexPointers` arrays, and O(m) for the answer.

Trade-offs

Pros

Optimal time complexity, linear in the input size and value range.
Very fast for the given constraints.

Cons

Slightly more complex to implement due to the need to manage multiple pointers.
Uses more memory if the maximum value W is very large, though it's fine for the given constraints.

Solutions

Solution

1class Solution {2public3  long[] unmarkedSumArray(int[] nums, int[][] queries) {4    int n = nums.length;5    long s = Arrays.stream(nums).asLongStream().sum();6    boolean[] mark = new boolean[n];7    int[][] arr = new int[n][0];8    for (int i = 0; i < n; ++i) {9      arr[i] = new int[]{nums[i], i};10    }11    Arrays.sort(arr, (a, b)->a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]);12    int m = queries.length;13    long[] ans = new long[m];14    for (int i = 0, j = 0; i < m; ++i) {15      int index = queries[i][0], k = queries[i][1];16      if (!mark[index]) {17        mark[index] = true;18        s -= nums[index];19      }20      for (; k > 0 && j < n; ++j) {21        if (!mark[arr[j][1]]) {22          mark[arr[j][1]] = true;23          s -= arr[j][0];24          --k;25        }26      }27      ans[i] = s;28    }29    return ans;30  }31}

1class Solution {2public:3  vector<long long> unmarkedSumArray(vector<int> &nums,4                                     vector<vector<int>> &queries) {5    int n = nums.size();6    long long s = accumulate(nums.begin(), nums.end(), 0LL);7    vector<bool> mark(n);8    vector<pair<int, int>> arr;9    for (int i = 0; i < n; ++i) {10      arr.emplace_back(nums[i], i);11    }12    sort(arr.begin(), arr.end());13    vector<long long> ans;14    int m = queries.size();15    for (int i = 0, j = 0; i < m; ++i) {16      int index = queries[i][0], k = queries[i][1];17      if (!mark[index]) {18        mark[index] = true;19        s -= nums[index];20      }21      for (; k && j < n; ++j) {22        if (!mark[arr[j].second]) {23          mark[arr[j].second] = true;24          s -= arr[j].first;25          --k;26        }27      }28      ans.push_back(s);29    }30    return ans;31  }32};

1class Solution:2    def unmarkedSumArray(self, nums: List[int], queries: List[List[int]]) -> List[int]: n = len(nums) s = sum(nums) mark = [False] * n arr = sorted((x, i) for i, x in enumerate(nums)) j = 0 ans = [] for index, k in queries: if not mark[index]: mark[index] = True s -= nums[index] while k and j < n: if not mark[arr[j][1]]: mark[arr[j][1]] = True s -= arr[j][0] k -= 1 j += 1 ans . append(s) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem

Solution

1class Solution {2public3  long[] unmarkedSumArray(int[] nums, int[][] queries) {4    int n = nums.length;5    long s = Arrays.stream(nums).asLongStream().sum();6    boolean[] mark = new boolean[n];7    int[][] arr = new int[n][0];8    for (int i = 0; i < n; ++i) {9      arr[i] = new int[]{nums[i], i};10    }11    Arrays.sort(arr, (a, b)->a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]);12    int m = queries.length;13    long[] ans = new long[m];14    for (int i = 0, j = 0; i < m; ++i) {15      int index = queries[i][0], k = queries[i][1];16      if (!mark[index]) {17        mark[index] = true;18        s -= nums[index];19      }20      for (; k > 0 && j < n; ++j) {21        if (!mark[arr[j][1]]) {22          mark[arr[j][1]] = true;23          s -= arr[j][0];24          --k;25        }26      }27      ans[i] = s;28    }29    return ans;30  }31}

1class Solution {2public:3  vector<long long> unmarkedSumArray(vector<int> &nums,4                                     vector<vector<int>> &queries) {5    int n = nums.size();6    long long s = accumulate(nums.begin(), nums.end(), 0LL);7    vector<bool> mark(n);8    vector<pair<int, int>> arr;9    for (int i = 0; i < n; ++i) {10      arr.emplace_back(nums[i], i);11    }12    sort(arr.begin(), arr.end());13    vector<long long> ans;14    int m = queries.size();15    for (int i = 0, j = 0; i < m; ++i) {16      int index = queries[i][0], k = queries[i][1];17      if (!mark[index]) {18        mark[index] = true;19        s -= nums[index];20      }21      for (; k && j < n; ++j) {22        if (!mark[arr[j].second]) {23          mark[arr[j].second] = true;24          s -= arr[j].first;25          --k;26        }27      }28      ans.push_back(s);29    }30    return ans;31  }32};

1class Solution:2    def unmarkedSumArray(self, nums: List[int], queries: List[List[int]]) -> List[int]: n = len(nums) s = sum(nums) mark = [False] * n arr = sorted((x, i) for i, x in enumerate(nums)) j = 0 ans = [] for index, k in queries: if not mark[index]: mark[index] = True s -= nums[index] while k and j < n: if not mark[arr[j][1]]: mark[arr[j][1]] = True s -= arr[j][0] k -= 1 j += 1 ans . append(s) return ans

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.