3318. Find X-Sum of All K-Long Subarrays I

EasyArrayHash TableSliding WindowHeap (Priority Queue)

Problem Description

You are given an array nums of n integers and two integers k and x.

The x-sum of an array is calculated by the following procedure:

Count the occurrences of all elements in the array.
Keep only the occurrences of the top x most frequent elements. If two elements have the same number of occurrences, the element with the bigger value is considered more frequent.
Calculate the sum of the resulting array.

Note that if an array has less than x distinct elements, its x-sum is the sum of the array.

Return an integer array answer of length n - k + 1 where answer[i] is the x-sum of the subarray nums[i..i + k - 1].

Intuition

To solve this problem efficiently, the goal is to compute the x-sum for each subarray of length k within the main array nums. Due to overlapping windows as we slide through nums, we need to maintain a structure that dynamically tracks frequency counts and efficiently provides the frequent elements for sum calculation.

Here's how we approach this:

Frequency Counting: Use a hash table to track the number of occurrences of each element within the current window of size k.
Maintain Top Frequent Elements: Utilize two ordered sets l and r.
- l stores the top x elements with the highest occurrences, and it's used to calculate the x-sum.
- r stores the rest and helps dynamically find new top elements when the window slides.
Dynamic Updates:
- When handling each element, increment its count, and adjust the position of elements between l and r to maintain the correct top x elements.
- Update the sum s by adding elements from l.
Sliding Window: Adjust the window by removing the element that slides out and updating counts and sets accordingly. This avoids rescanning the entire window, leveraging the overlaps to maintain efficiency.

This approach efficiently handles each window shift with updates in logarithmic time concerning the number of elements, making it suitable for larger arrays.

Learn more about Sliding Window and Heap (Priority Queue) patterns.

Solution Approach

The solution employs a combination of hash table (dictionary) and ordered sets (SortedList from sortedcontainers), along with a sliding window technique to efficiently compute the x-sum for each subarray of length k. Here's how this is implemented:

Data Structures:
- cnt: A Counter object from Python's collections to keep track of occurrences of each element within the current window.
- l and r: Two SortedList structures to maintain the elements based on frequency and value. l holds the top x elements, and r holds the rest.
Initialization:
- Initialize the sliding window and the associated data structures. Set s to 0 to accumulate the x-sum for each window.
- ans, an array, is initialized to store the results for each window starting point.
Functions:
- add(v): Adjust the sets when an element v is added. Increment its count, and adjust its position between l and r based on its frequency.
- remove(v): Adjust the sets when an element v is removed. Decrease its count and reposition it accordingly.
Sliding Window:
- For every element in nums, increment its count and adjust its placement using add(v).
- For each valid window (once it has k elements), ensure l has exactly x most frequent elements. Move elements between l and r as needed to maintain this constraint and update the sum s.
- Record s in ans at position j, where j is the start index of the current window.
- Slide the window by removing the left-most element and updating structures using remove(nums[j]).
Complexity:
- The approach achieves an efficient time complexity of O(n \times \log k), where n is the length of nums, due to the operations on the ordered sets.
- Space complexity is O(n), mainly driven by the storage requirements of the frequency counter and auxiliary lists.

This method effectively handles the complexities of dynamic subarray queries with respect to frequency and resultant sums, optimizing performance through organized data structures and careful maintenance of state between window shifts.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's consider a small example to illustrate the solution approach.

Suppose we have the following inputs:

nums = [1, 2, 2, 1, 3, 3, 3]
k = 4
x = 2

We aim to compute the x-sum for each subarray of length k.

Initial Setup:

Define l and r as empty SortedList structures.
Use a Counter to track the frequencies.

First Window (subarray nums[0..3] = [1, 2, 2, 1]):

Count frequencies: 1 occurs twice, 2 occurs twice.
l should contain the top 2 most frequent elements. Here both 1 and 2 have the same frequency. The element 2 is larger, so l = {2, 1}.
Calculate the sum of elements in l: x-sum = 2 + 1 = 3.
Record the result for this window: answer[0] = 3.

Slide the Window:

Remove nums[0] = 1.
Update frequencies: 1 has a reduced frequency.
Adjust sets l and r accordingly.

Second Window (subarray nums[1..4] = [2, 2, 1, 3]):

Update frequency counts with the new element 3.
Adjust l and r: the top two elements by frequency are [2] then [1 or 3] based on value, so choose [2, 3].
x-sum evaluated here is 2 + 3 = 5.
Record this x-sum: answer[1] = 5.

Continue Sliding:

Repeat the process for the next window nums[2..5] and nums[3..6], updating l, r, and the Counter.

Final Output:

After processing all such windows, the answer array forms: [3, 5, 6, 7], where each element corresponds to the x-sum calculation for the sliding windows.

Using this walkthrough, the solution efficiently calculates the x-sum for each sliding window length of k, handling updates by appropriately managing the ordered set structures and adjusting based on updated frequencies.

Solution Implementation

1from typing import List
2from collections import Counter
3from sortedcontainers import SortedList
4
5class Solution:
6    def findXSum(self, nums: List[int], k: int, x: int) -> List[int]:
7        # Function to add an integer to the sorted sets
8        def add(value: int):
9            # If the counter for the value is zero, ignore this request
10            if count[value] == 0:
11                return
12          
13            pair = (count[value], value)  # Create a tuple with count and value
14
15            # If there's enough elements in 'l' and 'pair' is greater than
16            # the smallest in 'l', update the sum and add to 'l'
17            if left and pair > left[0]:  
18                nonlocal sum_x
19                sum_x += pair[0] * pair[1]
20                left.add(pair)
21            else:
22                # Else add to 'r'
23                right.add(pair)
24
25        # Function to remove an integer from the sorted sets
26        def remove(value: int):
27            # If the counter for the value is zero, ignore this request
28            if count[value] == 0:
29                return
30          
31            pair = (count[value], value)  # Create a tuple with count and value
32
33            # Remove from 'left' if present, otherwise remove from 'right'
34            if pair in left:
35                nonlocal sum_x
36                sum_x -= pair[0] * pair[1]
37                left.remove(pair)
38            else:
39                right.remove(pair)
40
41        left = SortedList()   # Elements required for the answer
42        right = SortedList()  # Remaining elements
43        count = Counter()     # Keep track of occurrences of each number
44        sum_x = 0             # Sum of elements in 'left'
45        n = len(nums)         # Length of the input list
46        result = [0] * (n - k + 1)  # Array to store results
47
48        # Iterate over numbers in the list
49        for i, value in enumerate(nums):
50            remove(value)           # Remove current value's contribution
51            count[value] += 1       # Increase the count
52            add(value)              # Add current value's new contribution
53
54            j = i - k + 1
55            if j < 0:
56                continue
57
58            # Balance the 'left' list to maintain exactly 'x' elements
59            while right and len(left) < x:
60                pair = right.pop()
61                left.add(pair)
62                sum_x += pair[0] * pair[1]
63
64            while len(left) > x:
65                pair = left.pop(0)
66                sum_x -= pair[0] * pair[1]
67                right.add(pair)
68
69            # Store the calculated sum in result
70            result[j] = sum_x
71
72            # Prepare for the next window
73            remove(nums[j])
74            count[nums[j]] -= 1
75            add(nums[j])
76
77        return result
78

1import java.util.HashMap;
2import java.util.Map;
3import java.util.TreeSet;
4
5class Solution {
6    // TreeSets to manage two different sets of integer arrays with a custom comparator
7    private TreeSet<int[]> left = new TreeSet<>((a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]);
8    private TreeSet<int[]> right = new TreeSet<>(left.comparator());
9    private Map<Integer, Integer> countMap = new HashMap<>();
10    private int sum;
11
12    /**
13     * Find the x-sum of sliding window subarrays of size k in the array nums.
14     * @param nums The input array of integers.
15     * @param k The size of the sliding window.
16     * @param x The number of smallest elements to sum up in each window.
17     * @return An array of sums for each window position.
18     */
19    public int[] findXSum(int[] nums, int k, int x) {
20        int n = nums.length;
21        int[] result = new int[n - k + 1]; // Result array to store the sums for each window.
22
23        // Loop through each element in the input array.
24        for (int i = 0; i < n; ++i) {
25            int value = nums[i];
26            remove(value); // Remove current element from data structures.
27            countMap.merge(value, 1, Integer::sum); // Increment the count of current element.
28            add(value); // Add current element back to data structures.
29
30            int windowStart = i - k + 1;
31            if (windowStart < 0) {
32                continue; // Skip if the window is not yet fully overlapping the input array.
33            }
34
35            // Adjust the left set to ensure it contains exactly x smallest elements.
36            while (!right.isEmpty() && left.size() < x) {
37                var elem = right.pollLast(); // Move largest from right to left.
38                sum += elem[0] * elem[1];
39                left.add(elem);
40            }
41
42            // Remove excess elements from left set until it only contains x elements.
43            while (left.size() > x) {
44                var elem = left.pollFirst(); // Move smallest from left to right.
45                sum -= elem[0] * elem[1];
46                right.add(elem);
47            }
48
49            result[windowStart] = sum; // Store the calculated sum for the current window.
50
51            remove(nums[windowStart]); // Prepare to slide the window.
52            countMap.merge(nums[windowStart], -1, Integer::sum); // Decrease the count of outgoing element.
53            add(nums[windowStart]); // Re-add to ensure correctness for next window.
54        }
55
56        return result; // Return the resultant x-sums for all windows.
57    }
58
59    // Remove an element from the data structures
60    private void remove(int value) {
61        if (!countMap.containsKey(value)) {
62            return;
63        }
64        var entry = new int[] {countMap.get(value), value};
65        if (left.contains(entry)) {
66            left.remove(entry);
67            sum -= entry[0] * entry[1];
68        } else {
69            right.remove(entry);
70        }
71    }
72
73    // Add an element back to the correct data structures
74    private void add(int value) {
75        if (!countMap.containsKey(value)) {
76            return;
77        }
78        var entry = new int[] {countMap.get(value), value};
79        if (!left.isEmpty() && left.comparator().compare(left.first(), entry) < 0) {
80            left.add(entry);
81            sum += entry[0] * entry[1];
82        } else {
83            right.add(entry);
84        }
85    }
86}
87

1class Solution {
2public:
3    std::vector<int> findXSum(std::vector<int>& nums, int k, int x) {
4        using PairInt = std::pair<int, int>; // Define a pair type for frequency and value
5        std::set<PairInt> leftSet, rightSet; // Sets to maintain current active elements
6        int sum = 0; // To hold the sum of top 'x' frequent elements
7        std::unordered_map<int, int> frequency; // Map to count frequency of elements
8
9        // Lambda function to add element 'v' into the appropriate set
10        auto add = [&](int v) {
11            if (frequency[v] == 0) return; // Skip if element has zero frequency
12
13            PairInt p = {frequency[v], v}; // Create a pair of frequency and element value
14
15            if (!leftSet.empty() && p > *leftSet.begin()) { // Check position based on frequency and current smallest in leftSet
16                sum += p.first * p.second; // Add product of frequency and value to sum
17                leftSet.insert(p); // Insert into leftSet
18            } else {
19                rightSet.insert(p); // Otherwise, insert into rightSet
20            }
21        };
22
23        // Lambda function to remove element 'v' from the appropriate set
24        auto remove = [&](int v) {
25            if (frequency[v] == 0) return; // Skip if element has zero frequency
26
27            PairInt p = {frequency[v], v}; // Create a pair of frequency and value
28            auto it = leftSet.find(p);
29
30            if (it != leftSet.end()) { // If found in leftSet, remove and adjust sum
31                sum -= p.first * p.second;
32                leftSet.erase(it);
33            } else {
34                rightSet.erase(p); // Else, remove from rightSet
35            }
36        };
37
38        std::vector<int> result; // Result vector to store answers
39
40        for (int i = 0; i < nums.size(); ++i) {
41            remove(nums[i]); // Remove current element from its previous position
42            ++frequency[nums[i]]; // Increase frequency of current element
43            add(nums[i]); // Add current element to the sets
44
45            int j = i - k + 1; // Start of the current window
46
47            if (j < 0) continue; // Skip until full window is formed
48
49            // Balance sets if needed, ensuring leftSet has exactly 'x' elements
50            while (!rightSet.empty() && leftSet.size() < x) {
51                PairInt p = *rightSet.rbegin(); // Get the last element in rightSet
52                sum += p.first * p.second; // Add it to sum
53                rightSet.erase(p);
54                leftSet.insert(p);
55            }
56            while (leftSet.size() > x) {
57                PairInt p = *leftSet.begin(); // Get the smallest element in leftSet
58                sum -= p.first * p.second; // Subtract it from sum
59                leftSet.erase(p);
60                rightSet.insert(p);
61            }
62            result.push_back(sum); // Add current sum to result
63
64            remove(nums[j]); // Remove element that's sliding out of the window
65            --frequency[nums[j]]; // Decrease its frequency
66            add(nums[j]); // Re-add it to the sets
67        }
68
69        return result; // Return the resulting sums
70    }
71};
72

1type PairInt = [number, number]; // Define a tuple for frequency and value
2
3// Sets to maintain current active elements
4let leftSet: Set<PairInt> = new Set();
5let rightSet: Set<PairInt> = new Set();
6
7let sum = 0; // To hold the sum of top 'x' frequent elements
8
9// Map to count frequency of elements
10let frequency: Map<number, number> = new Map();
11
12// Function to add element 'v' into the appropriate set
13const add = (v: number) => {
14    if ((frequency.get(v) || 0) === 0) return; // Skip if element has zero frequency
15
16    const p: PairInt = [frequency.get(v) || 0, v]; // Create a pair of frequency and element value
17
18    if (leftSet.size > 0 && comparePairs(p, getMin(leftSet))) { // Check position based on frequency and current smallest in leftSet
19        sum += p[0] * p[1]; // Add product of frequency and value to sum
20        leftSet.add(p); // Insert into leftSet
21    } else {
22        rightSet.add(p); // Otherwise, insert into rightSet
23    }
24};
25
26// Function to remove element 'v' from the appropriate set
27const remove = (v: number) => {
28    if ((frequency.get(v) || 0) === 0) return; // Skip if element has zero frequency
29
30    const p: PairInt = [frequency.get(v) || 0, v]; // Create a pair of frequency and value
31
32    if (leftSet.has(p)) { // If found in leftSet, remove and adjust sum
33        sum -= p[0] * p[1];
34        leftSet.delete(p);
35    } else {
36        rightSet.delete(p); // Else, remove from rightSet
37    }
38};
39
40// Compare two pairs (frequency, value)
41const comparePairs = (a: PairInt, b: PairInt) => a[0] > b[0] || (a[0] === b[0] && a[1] > b[1]);
42
43// Get the smallest element in a set (this implementation assumes pairs are sorted internally)
44const getMin = (s: Set<PairInt>) => {
45    let min = Array.from(s)[0];
46    s.forEach(e => {
47        if (comparePairs(min, e)) min = e;
48    });
49    return min;
50};
51
52// Get the largest element in a set (this implementation assumes pairs are sorted internally)
53const getMax = (s: Set<PairInt>) => {
54    let max = Array.from(s)[0];
55    s.forEach(e => {
56        if (comparePairs(e, max)) max = e;
57    });
58    return max;
59};
60
61// Function to find the sum of the 'x' most frequent elements in each sliding window of size 'k'
62function findXSum(nums: number[], k: number, x: number): number[] {
63    let result: number[] = []; // Result array to store answers
64
65    for (let i = 0; i < nums.length; ++i) {
66        remove(nums[i]); // Remove current element from its previous position
67        frequency.set(nums[i], (frequency.get(nums[i]) || 0) + 1); // Increase frequency of current element
68        add(nums[i]); // Add current element to the sets
69
70        const j = i - k + 1; // Start of the current window
71
72        if (j < 0) continue; // Skip until full window is formed
73
74        // Balance sets if needed, ensuring leftSet has exactly 'x' elements
75        while (rightSet.size > 0 && leftSet.size < x) {
76            const p = getMax(rightSet); // Get the last element in rightSet
77            sum += p[0] * p[1]; // Add it to sum
78            rightSet.delete(p);
79            leftSet.add(p);
80        }
81        while (leftSet.size > x) {
82            const p = getMin(leftSet); // Get the smallest element in leftSet
83            sum -= p[0] * p[1]; // Subtract it from sum
84            leftSet.delete(p);
85            rightSet.add(p);
86        }
87        result.push(sum); // Add current sum to result
88
89        remove(nums[j]); // Remove element that's sliding out of the window
90        frequency.set(nums[j], (frequency.get(nums[j]) || 0) - 1); // Decrease its frequency
91        add(nums[j]); // Re-add it to the sets
92    }
93
94    return result; // Return the resulting sums
95}
96

Time and Space Complexity

The code's time complexity is O(n log n). This is due to the fact that the operations involving SortedList, such as adding or removing elements (add and remove), and balancing elements between l and r lists, can be done in O(log n). The loop iterates over the entire list of nums, which is O(n).

The space complexity is O(n). This is primarily because the SortedList instances l and r as well as the Counter object can collectively hold up to O(n) elements in the worst case.

Learn more about how to find time and space complexity quickly.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan: