3321. Find X-Sum of All K-Long Subarrays II
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, we use a combination of a hash table and ordered sets:
-
Hash Table
cnt
: This keeps track of the frequency of each element within the sliding window. -
Ordered Sets
l
andr
: These are used to store the elements based on their frequencies. The ordered setl
stores the topx
elements with the highest occurrences. The ordered setr
stores the remaining elements. -
Sum
s
: This variable maintains the sum of the elements currently inl
. As we slide the window, we updatel
,r
, ands
.
The process involves:
- Initializing the window with the first
k
elements. - Traversing the array while maintaining the
x
most frequent elements inl
. - Moving elements between
l
andr
to keepl
exactly with thex
highest frequency elements. - Calculating the
x-sum
for each valid window configuration and storing it inanswer
.
This approach ensures efficient updates as we slide the window across nums
, taking advantage of the logarithmic operations allowed by ordered sets and the constant-time operations provided by hash tables.
Learn more about Sliding Window and Heap (Priority Queue) patterns.
Solution Approach
The solution uses a combination of the hash table and ordered sets to efficiently compute the x-sum over sliding windows of size k
within the given array nums
.
Key Components:
-
Hash Table (
cnt
): This keeps track of the count of each element in the current window. It's crucial for knowing when to update the frequency of an element as it enters or exits the sliding window. -
Ordered Sets (
l
andr
): These are used to maintain the order of elements by frequency.l
always contains the topx
most frequent elements or all elements if fewer exist.r
holds the rest of the elements in the window. We rely on these data structures to quickly adjust which elements are considered the most frequent. -
Sliding Window Technique: By using a sliding window, we efficiently process each subarray of length
k
to determine their x-sums.
Algorithm Steps:
-
Initialize Variables: Start with two empty
SortedList
structures,l
andr
, for ordered frequency management, aCounter
for element counting, the sums
initialized to zero, and ananswer
array for storing results. -
Element Addition and Removal:
- Add: When an element is added to the window, update its count. Depending on its frequency, decide whether it should be in
l
orr
, adjustings
accordingly. - Remove: When an element exits the window, decrease its count and update the ordered sets and sum
s
similarly.
- Add: When an element is added to the window, update its count. Depending on its frequency, decide whether it should be in
-
Sliding and Balance:
- Traverse each element of
nums
. For each element that enters the window, adjustl
andr
. - Adjust the balance between
l
andr
to ensurel
always contains the topx
frequency elements. - If
l
has fewer thanx
elements andr
is not empty, transfer the highest elements fromr
tol
. - If
l
exceeds the limitx
, transfer the lowest elements froml
tor
.
- Traverse each element of
-
Calculate and Store: Once the window is correctly adjusted, store the current sum
s
as the answer for the position corresponding to the left end of the window.
This approach efficiently maintains the desired frequency constraints due to the logarithmic operations afforded by SortedList
and constant-time operations from Counter
, achieving a time complexity of O(n \times \log k)
and a space complexity of O(n)
.
Ready to land your dream job?
Unlock your dream job with a 2-minute evaluator for a personalized learning plan!
Start EvaluatorExample Walkthrough
Let's consider a simple example to illustrate the solution approach:
Suppose we have an array nums = [4, 1, 2, 2, 3, 4]
, and we choose k = 3
and x = 2
. We need to find the x-sum
for each sliding window of size k
.
Step 1: Initialize
- Start with empty structures:
cnt
(a Counter): {}l
(SortedList): []r
(SortedList): []s
(sum of elements inl
): 0answer
: []
Step 2: Process First Window (nums[0..2] = [4, 1, 2]
)
-
Add elements into the window:
- Add
4
:cnt = {4: 1}
,l = [4]
,s = 4
- Add
1
:cnt = {4: 1, 1: 1}
,l = [4, 1]
,s = 5
(sincex=2
, move 1 tor
) - Add
2
:cnt = {4: 1, 1: 1, 2: 1}
,r = [1, 2]
- Add
-
Adjust
l
andr
:l
contains4
, the top frequent element in the window
-
Calculate
x-sum
:s = 4
-
Append
x-sum
toanswer
:answer = [4]
Step 3: Slide the Window
-
Slide by one element to
nums[1..3] = [1, 2, 2]
-
Remove
4
:cnt = {1: 1, 2: 1}
,l = []
,r = [1, 2]
(transfer2
tol
) -
Add
2
:cnt = {1: 1, 2: 2}
,l = [2]
,s = 2
-
Adjust
l
andr
:l
contains2
, the element with highest frequency
-
Calculate
x-sum
:s = 2
-
Update
answer
:answer = [4, 2]
Step 4: Continue Sliding
-
Slide to
nums[2..4] = [2, 2, 3]
-
Remove
1
:cnt = {2: 2, 3: 1}
,l = [2]
,s = 2
-
Add
3
:cnt = {2: 2, 3: 1}
,r = [3]
-
Adjust
l
andr
:l
already has the most frequent (2
)
-
Calculate
x-sum
:s = 2
-
Update
answer
:answer = [4, 2, 2]
Step 5: Final Slide
-
Slide to
nums[3..5] = [2, 3, 4]
-
Remove
2
:cnt = {2: 1, 3: 1, 4: 1}
,l = [3, 4]
,s = 7
-
Add
4
:cnt = {2: 1, 3: 1, 4: 2}
,l = [4]
,s = 4
-
Calculate
x-sum
:s = 4
-
Final
answer
:answer = [4, 2, 2, 4]
This walkthrough demonstrates how the solution maintains the top x
frequency elements efficiently using sliding windows, hash tables, and ordered sets.
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 # Helper function to add an element to the appropriate set
8 def add(value: int):
9 if count[value] == 0:
10 return
11 element = (count[value], value)
12 # Check if the element should be in the 'left' set
13 if left and element > left[0]:
14 nonlocal current_sum
15 # Update the sum for the 'left' set
16 current_sum += element[0] * element[1]
17 left.add(element)
18 else:
19 # Add the element to the 'right' set
20 right.add(element)
21
22 # Helper function to remove an element from the appropriate set
23 def remove(value: int):
24 if count[value] == 0:
25 return
26 element = (count[value], value)
27 if element in left:
28 nonlocal current_sum
29 # Update the sum by removing the element's contribution
30 current_sum -= element[0] * element[1]
31 left.remove(element)
32 else:
33 # Remove the element from the 'right' set
34 right.remove(element)
35
36 left = SortedList() # Set for storing the smallest elements by some criterion
37 right = SortedList() # Set for storing other elements
38 count = Counter() # Counter to keep track of elements' frequencies
39 current_sum = 0 # Variable to maintain the sum of the 'left' set
40 n = len(nums)
41 ans = [0] * (n - k + 1) # Result list to store sums of the valid subarrays
42
43 # Iterate through each element in nums
44 for i, value in enumerate(nums):
45 remove(value) # Attempt to remove the element from consideration for both sets
46 count[value] += 1 # Increment the count for this element
47 add(value) # Re-add the element with updated count to the appropriate set
48
49 j = i - k + 1 # Calculate the starting index of the current window
50 if j < 0:
51 continue # Skip if the window isn't fully formed yet
52
53 # Balance the sets: Move elements from 'right' to 'left' if needed
54 while right and len(left) < x:
55 element = right.pop() # Get an element
56 left.add(element) # Add to 'left'
57 current_sum += element[0] * element[1] # Add to current_sum
58
59 # Ensure 'left' has exactly 'x' elements
60 while len(left) > x:
61 element = left.pop(0) # Remove the smallest element from 'left'
62 current_sum -= element[0] * element[1] # Subtract from current_sum
63 right.add(element) # Add to 'right'
64
65 ans[j] = current_sum # Store the current sum of 'x' largest occurring elements
66
67 # Prepare for the next window by updating counts and sets
68 remove(nums[j])
69 count[nums[j]] -= 1
70 add(nums[j])
71
72 return ans
73
1import java.util.HashMap;
2import java.util.Map;
3import java.util.TreeSet;
4
5class Solution {
6 // Comparator to sort integer arrays based on their first element,
7 // and by second element if first elements are equal
8 private TreeSet<int[]> leftSet = new TreeSet<>((a, b) -> a[0] == b[0] ? a[1] - b[1] : a[0] - b[0]);
9 // TreeSet with the same comparator as leftSet
10 private TreeSet<int[]> rightSet = new TreeSet<>(leftSet.comparator());
11 // Map to keep track of the frequency of each integer
12 private Map<Integer, Integer> countMap = new HashMap<>();
13 // Sum of integer products within the specified range
14 private long sum;
15
16 public long[] findXSum(int[] nums, int k, int x) {
17 int n = nums.length;
18 long[] result = new long[n - k + 1];
19
20 for (int i = 0; i < n; ++i) {
21 int value = nums[i];
22
23 // Remove current value from the sets and update frequency
24 remove(value);
25 countMap.merge(value, 1, Integer::sum); // Add 1 to the count of the value
26 add(value); // Add the value with updated frequency
27
28 int j = i - k + 1;
29 if (j < 0) {
30 continue;
31 }
32
33 // Balance the size of leftSet until it contains x elements
34 while (!rightSet.isEmpty() && leftSet.size() < x) {
35 var pair = rightSet.pollLast(); // Remove largest from rightSet
36 sum += 1L * pair[0] * pair[1]; // Add to sum using elements of pair
37 leftSet.add(pair); // Add pair to leftSet
38 }
39
40 // Ensure leftSet doesn't exceed x elements
41 while (leftSet.size() > x) {
42 var pair = leftSet.pollFirst(); // Remove smallest from leftSet
43 sum -= 1L * pair[0] * pair[1]; // Subtract from sum
44 rightSet.add(pair); // Add pair to rightSet
45 }
46
47 // Record computed sum in the answer array
48 result[j] = sum;
49
50 // Update structures by removing the element that has slid out of the window
51 remove(nums[j]);
52 countMap.merge(nums[j], -1, Integer::sum); // Decrease count of the value
53 add(nums[j]);
54 }
55 return result;
56 }
57
58 private void remove(int value) {
59 // If the value is not in the map, nothing to remove
60 if (!countMap.containsKey(value)) {
61 return;
62 }
63
64 var pair = new int[]{countMap.get(value), value};
65
66 if (leftSet.contains(pair)) {
67 // Remove from leftSet if present, and adjust the sum
68 leftSet.remove(pair);
69 sum -= 1L * pair[0] * pair[1];
70 } else {
71 // Remove from rightSet
72 rightSet.remove(pair);
73 }
74 }
75
76 private void add(int value) {
77 // If the value is not in the map, nothing to add
78 if (!countMap.containsKey(value)) {
79 return;
80 }
81
82 var pair = new int[]{countMap.get(value), value};
83
84 if (!leftSet.isEmpty() && leftSet.comparator().compare(leftSet.first(), pair) < 0) {
85 // Add to leftSet if it should belong there, and adjust the sum
86 leftSet.add(pair);
87 sum += 1L * pair[0] * pair[1];
88 } else {
89 // Otherwise, add to rightSet
90 rightSet.add(pair);
91 }
92 }
93}
94
1#include <vector>
2#include <unordered_map>
3#include <set>
4
5class Solution {
6public:
7 // Function to find the x-sum of subarrays of size k in a list of integers.
8 std::vector<long long> findXSum(std::vector<int>& nums, int k, int x) {
9 using pii = std::pair<int, int>; // Define a pair used for set operations.
10 std::set<pii> leftSet, rightSet; // Two sets to manage elements based on their contribution to the sum.
11 long long currentSum = 0;
12 std::unordered_map<int, int> countMap; // Map to store the frequency of elements in the current window.
13
14 // Lambda function to add an element into the appropriate set.
15 auto addElement = [&](int value) {
16 if (countMap[value] == 0) return;
17 pii element = {countMap[value], value};
18 if (!leftSet.empty() && element > *leftSet.begin()) {
19 currentSum += 1LL * element.first * element.second;
20 leftSet.insert(element);
21 } else {
22 rightSet.insert(element);
23 }
24 };
25
26 // Lambda function to remove an element from the appropriate set.
27 auto removeElement = [&](int value) {
28 if (countMap[value] == 0) return;
29 pii element = {countMap[value], value};
30 auto it = leftSet.find(element);
31 if (it != leftSet.end()) {
32 currentSum -= 1LL * element.first * element.second;
33 leftSet.erase(it);
34 } else {
35 rightSet.erase(element);
36 }
37 };
38
39 std::vector<long long> result;
40
41 for (int i = 0; i < nums.size(); ++i) {
42 // Update sets and count for the current element.
43 removeElement(nums[i]);
44 ++countMap[nums[i]];
45 addElement(nums[i]);
46
47 int j = i - k + 1;
48 if (j < 0) continue; // Continue if the window is not yet of size k.
49
50 // Balance the sets to maintain exactly x elements in the left set.
51 while (!rightSet.empty() && leftSet.size() < x) {
52 pii element = *rightSet.rbegin();
53 currentSum += 1LL * element.first * element.second;
54 rightSet.erase(element);
55 leftSet.insert(element);
56 }
57 while (leftSet.size() > x) {
58 pii element = *leftSet.begin();
59 currentSum -= 1LL * element.first * element.second;
60 leftSet.erase(element);
61 rightSet.insert(element);
62 }
63 result.push_back(currentSum); // Store the current sum in the results.
64
65 // Prepare for the next iteration by removing the element going out of the window.
66 removeElement(nums[j]);
67 --countMap[nums[j]];
68 addElement(nums[j]);
69 }
70 return result; // Return the list of x-sums for subarrays.
71 }
72};
73
1// Import necessary modules
2type Pair = [number, number];
3
4let leftSet = new Set<Pair>(); // Set to manage elements part of the x-sum.
5let rightSet = new Set<Pair>(); // Set to manage other elements.
6let currentSum = 0; // Stores the current x-sum of elements.
7let countMap = new Map<number, number>(); // Map to track the frequency of elements in the current window.
8
9function findXSum(nums: number[], k: number, x: number): number[] {
10 // Lambda function to add an element into the appropriate set.
11 const addElement = (value: number) => {
12 const count = countMap.get(value) || 0;
13 if (count === 0) return;
14
15 const element: Pair = [count, value];
16 if (leftSet.size > 0 && compareElements(element, getSetFirstElement(leftSet))) {
17 currentSum += BigInt(element[0]) * BigInt(element[1]);
18 leftSet.add(element);
19 } else {
20 rightSet.add(element);
21 }
22 };
23
24 // Lambda function to remove an element from the appropriate set.
25 const removeElement = (value: number) => {
26 const count = countMap.get(value) || 0;
27 if (count === 0) return;
28
29 const element: Pair = [count, value];
30 if (leftSet.has(element)) {
31 currentSum -= BigInt(element[0]) * BigInt(element[1]);
32 leftSet.delete(element);
33 } else {
34 rightSet.delete(element);
35 }
36 };
37
38 // Initialize result array to store the x-sum values for subarrays of size k.
39 const result: number[] = [];
40
41 for (let i = 0; i < nums.length; ++i) {
42 // Update sets and count for the current element.
43 removeElement(nums[i]);
44 countMap.set(nums[i], (countMap.get(nums[i]) || 0) + 1);
45 addElement(nums[i]);
46
47 const j = i - k + 1;
48 if (j < 0) continue; // Skip if the window is not yet of size k.
49
50 // Balance the sets to maintain exactly x elements in the left set.
51 while (rightSet.size > 0 && leftSet.size < x) {
52 const element = getSetLastElement(rightSet);
53 currentSum += BigInt(element[0]) * BigInt(element[1]);
54 rightSet.delete(element);
55 leftSet.add(element);
56 }
57
58 while (leftSet.size > x) {
59 const element = getSetFirstElement(leftSet);
60 currentSum -= BigInt(element[0]) * BigInt(element[1]);
61 leftSet.delete(element);
62 rightSet.add(element);
63 }
64
65 result.push(Number(currentSum)); // Store the current sum in the results.
66
67 // Prepare for the next iteration by removing the element going out of the window.
68 removeElement(nums[j]);
69 countMap.set(nums[j], (countMap.get(nums[j]) || 0) - 1);
70 addElement(nums[j]);
71 }
72
73 return result; // Return the list of x-sums for subarrays.
74}
75
76// Helper function to compare two elements.
77function compareElements(a: Pair, b: Pair): boolean {
78 return a[0] > b[0] || (a[0] === b[0] && a[1] > b[1]);
79}
80
81// Helper function to get the first (minimum) element from a set.
82function getSetFirstElement(set: Set<Pair>): Pair {
83 return [...set][0]; // Deconstructs set to array and return first element
84}
85
86// Helper function to get the last (maximum) element from a set.
87function getSetLastElement(set: Set<Pair>): Pair {
88 return [...set].slice(-1)[0];
89}
90
Time and Space Complexity
The findXSum
function uses several operations like adding and removing elements from SortedList
, which is based on binary search trees. Each add or remove operation in a SortedList
has a complexity of O(log n)
.
The sliding window iterates through the entire array nums
, which has n
elements. For each element, the operations performed are:
-
Removing an element from
r
orl
(eachO(log x)
), adding the updated element back, and possibly adjusting which elements are inl
versusr
. The adjustment is done with binary operations on sets, effectively costingO(log x)
. -
In worst case scenarios, each move of the window entails toggling potentially all window elements, leading to moves among
l
andr
. The operations inside while loops (while r and len(l) < x:
andwhile len(l) > x:
) include operations on sorted data structures which altogether can addO(x log x)
complexity because the size ofl
could go up tox
.
Hence, the overall time complexity becomes O(n * log x)
due to each of the above operations happening for each element over the n
elements.
The space complexity is primarily dictated by the data structures SortedList
for l
and r
and the counter cnt
. Since l
and r
can grow to size x
, and cnt
manages the full window (up to approximately k
or n
), the space required is O(k + x)
(or O(n + x)
in the worst case when k >> x
).
Learn more about how to find time and space complexity quickly.
How many ways can you arrange the three letters A, B and C?
Recommended Readings
https algomonster s3 us east 2 amazonaws com cover_photos stack svg Sliding Window Maximum Monotonic Stack We have an array and a sliding window defined by a start index and an end index The sliding window moves from left of the array to right There are always k elements in
https algomonster s3 us east 2 amazonaws com cover_photos heap svg Priority Queue and Heap What is the relationship between priority queue and heap Priority Queue is an Abstract Data Type and Heap is the concrete data structure we use to implement a priority queue Priority Queue A priority queue
Coding Interview Patterns Your Personal Dijkstra's Algorithm to Landing Your Dream Job The goal of AlgoMonster is to help you get a job in the shortest amount of time possible in a data driven way We compiled datasets of tech interview problems and broke them down by patterns This way
Want a Structured Path to Master System Design Too? Don’t Miss This!