Facebook Pixel

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:

  1. Count the occurrences of all elements in the array.
  2. 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.
  3. 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:

  1. Hash Table cnt: This keeps track of the frequency of each element within the sliding window.

  2. Ordered Sets l and r: These are used to store the elements based on their frequencies. The ordered set l stores the top x elements with the highest occurrences. The ordered set r stores the remaining elements.

  3. Sum s: This variable maintains the sum of the elements currently in l. As we slide the window, we update l, r, and s.

The process involves:

  • Initializing the window with the first k elements.
  • Traversing the array while maintaining the x most frequent elements in l.
  • Moving elements between l and r to keep l exactly with the x highest frequency elements.
  • Calculating the x-sum for each valid window configuration and storing it in answer.

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:

  1. 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.

  2. Ordered Sets (l and r): These are used to maintain the order of elements by frequency. l always contains the top x 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.

  3. 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 and r, for ordered frequency management, a Counter for element counting, the sum s initialized to zero, and an answer 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 or r, adjusting s accordingly.
    • Remove: When an element exits the window, decrease its count and update the ordered sets and sum s similarly.
  • Sliding and Balance:

    • Traverse each element of nums. For each element that enters the window, adjust l and r.
    • Adjust the balance between l and r to ensure l always contains the top x frequency elements.
    • If l has fewer than x elements and r is not empty, transfer the highest elements from r to l.
    • If l exceeds the limit x, transfer the lowest elements from l to r.
  • 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 Evaluator

Example 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 in l): 0
    • answer: []

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 (since x=2, move 1 to r)
    • Add 2: cnt = {4: 1, 1: 1, 2: 1}, r = [1, 2]
  • Adjust l and r:

    • l contains 4, the top frequent element in the window
  • Calculate x-sum: s = 4

  • Append x-sum to answer: 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] (transfer 2 to l)

  • Add 2: cnt = {1: 1, 2: 2}, l = [2], s = 2

  • Adjust l and r:

    • l contains 2, 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 and r:

    • 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:

  1. Removing an element from r or l (each O(log x)), adding the updated element back, and possibly adjusting which elements are in l versus r. The adjustment is done with binary operations on sets, effectively costing O(log x).

  2. In worst case scenarios, each move of the window entails toggling potentially all window elements, leading to moves among l and r. The operations inside while loops (while r and len(l) < x: and while len(l) > x:) include operations on sorted data structures which altogether can add O(x log x) complexity because the size of l could go up to x.

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.


Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

How many ways can you arrange the three letters A, B and C?


Recommended Readings

Want a Structured Path to Master System Design Too? Don’t Miss This!


Load More