3080. Mark Elements on Array by Performing Queries

Description

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

Solutions

Solution 1: Sorting + Simulation

First, we calculate the sum $s$ of the array $nums$. We define an array $mark$ to indicate whether the elements in the array have been marked, initializing all elements as unmarked.

Then, we create an array $arr$, where each element is a tuple $(x, i)$, indicating that the $i$-th element in the array has a value of $x$. We sort the array $arr$ by the value of the elements. If the values are equal, we sort them in ascending order of the index.

Next, we traverse the array $queries$. For each query $[index, k]$, we first check whether the element at index $index$ has been marked. If it has not been marked, we mark it and subtract the value of the element at index $index$ from $s$. Then, we traverse the array $arr$. For each element $(x, i)$, if element $i$ has not been marked, we mark it and subtract the value $x$ corresponding to element $i$ from $s$, until $k$ is $0$ or the array $arr$ is fully traversed. Then, we add $s$ to the answer array.

After traversing all the queries, we get the answer array.

The time complexity is $O(n \times \log n)$, and the space complexity is $O(n)$. Here, $n$ is the length of the array $nums$.

Python3JavaC++GoTypeScript

class Solution:
    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

class Solution {
    public long[] unmarkedSumArray(int[] nums, int[][] queries) {
        int n = nums.length;
        long s = Arrays.stream(nums).asLongStream().sum();
        boolean[] mark = new boolean[n];
        int[][] arr = new int[n][0];
        for (int i = 0; i < n; ++i) {
            arr[i] = new int[]{nums[i], i};
        }
        Arrays.sort(arr, (a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]);
        int m = queries.length;
        long[] ans = new long[m];
        for (int i = 0, j = 0; i < m; ++i) {
            int index = queries[i][0], k = queries[i][1];
            if (!mark[index]) {
                mark[index] = true;
                s -= nums[index];
            }
            for (; k > 0 && j < n; ++j) {
                if (!mark[arr[j][1]]) {
                    mark[arr[j][1]] = true;
                    s -= arr[j][0];
                    --k;
                }
            }
            ans[i] = s;
        }
        return ans;
    }
}

class Solution {
public:
    vector<long long> unmarkedSumArray(vector<int>& nums, vector<vector<int>>& queries) {
        int n = nums.size();
        long long s = accumulate(nums.begin(), nums.end(), 0LL);
        vector<bool> mark(n);
        vector<pair<int, int>> arr;
        for (int i = 0; i < n; ++i) {
            arr.emplace_back(nums[i], i);
        }
        sort(arr.begin(), arr.end());
        vector<long long> ans;
        int m = queries.size();
        for (int i = 0, j = 0; i < m; ++i) {
            int index = queries[i][0], k = queries[i][1];
            if (!mark[index]) {
                mark[index] = true;
                s -= nums[index];
            }
            for (; k && j < n; ++j) {
                if (!mark[arr[j].second]) {
                    mark[arr[j].second] = true;
                    s -= arr[j].first;
                    --k;
                }
            }
            ans.push_back(s);
        }
        return ans;
    }
};

func unmarkedSumArray(nums []int, queries [][]int) []int64 {
    n := len(nums)
    var s int64
    for _, x := range nums {
        s += int64(x)
    }
    mark := make([]bool, n)
    arr := make([][2]int, 0, n)
    for i, x := range nums {
        arr = append(arr, [2]int{x, i})
    }
    sort.Slice(arr, func(i, j int) bool {
        if arr[i][0] == arr[j][0] {
            return arr[i][1] < arr[j][1]
        }
        return arr[i][0] < arr[j][0]
    })
    ans := make([]int64, len(queries))
    j := 0
    for i, q := range queries {
        index, k := q[0], q[1]
        if !mark[index] {
            mark[index] = true
            s -= int64(nums[index])
        }
        for ; k > 0 && j < n; j++ {
            if !mark[arr[j][1]] {
                mark[arr[j][1]] = true
                s -= int64(arr[j][0])
                k--
            }
        }
        ans[i] = s
    }
    return ans
}

function unmarkedSumArray(nums: number[], queries: number[][]): number[] {
    const n = nums.length;
    let s = nums.reduce((acc, x) => acc + x, 0);
    const mark: boolean[] = Array(n).fill(false);
    const arr = nums.map((x, i) => [x, i]);
    arr.sort((a, b) => (a[0] === b[0] ? a[1] - b[1] : a[0] - b[0]));
    let j = 0;
    const ans: number[] = [];
    for (let [index, k] of queries) {
        if (!mark[index]) {
            mark[index] = true;
            s -= nums[index];
        }
        for (; k && j < n; ++j) {
            if (!mark[arr[j][1]]) {
                mark[arr[j][1]] = true;
                s -= arr[j][0];
                --k;
            }
        }
        ans.push(s);
    }
    return ans;
}

3080. Mark Elements on Array by Performing Queries

Description

Solutions

Solution 1: Sorting + Simulation

Comments