LeetCode 3321 - Find X-Sum of All K-Long Subarrays II

LeetCode Problem 3321

Difficulty: 🔴 Hard
Topics: Array, Hash Table, Sliding Window, Heap (Priority Queue)

Solution

Problem Understanding

The problem asks us to compute a specialized sum, called the x-sum, over all subarrays of length k in a given integer array nums. Specifically, the x-sum is computed by first counting the occurrences of each element in the subarray, keeping only the top x most frequent elements (with ties broken in favor of the larger element), and summing their total contribution. The input consists of an array of integers nums, a subarray length k, and a parameter x controlling how many top elements to consider for the sum. The output is an array of length n - k + 1 containing the x-sum for each contiguous subarray of length k.

The constraints indicate that the array can be fairly large, up to 10^5 elements, and the integers can be very large, up to 10^9. This suggests that any solution that examines all elements naively in each subarray could be too slow. Edge cases include when x equals k, where all elements in the subarray contribute to the sum, or when the subarray has fewer than x distinct elements, in which case all elements are considered.

Key points to notice are that elements with equal frequency are ordered by value, and that frequency counting is required for every sliding window of size k.

Approaches

The brute-force approach is straightforward. For each subarray of length k, compute the frequency of each element, sort the elements by frequency and value to select the top x, and sum their total. While correct, this approach is extremely inefficient because frequency counting and sorting must be repeated for every subarray, leading to a complexity of roughly O(n * k log k), which is too slow for n up to 10^5.

The optimal approach leverages the sliding window technique combined with a balanced multiset or heap to maintain the frequencies of elements efficiently as the window moves. The key insight is that when the window slides by one position, only one element leaves and one element enters. By updating the frequency counts incrementally and maintaining a heap or sorted structure for the top x elements, we can compute each window's x-sum in logarithmic time relative to the number of distinct elements in the window rather than the full subarray length.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n * k log k)	O(k)	Count and sort frequencies for each subarray
Optimal	O(n log k)	O(k)	Sliding window with balanced multiset or heap for frequencies

Algorithm Walkthrough

1. Initialize a frequency map to count occurrences of each element in the current window.
2. Use a max-heap (priority queue) or sorted structure to maintain elements sorted first by frequency descending, then by value descending.
3. Populate the initial window of size k by counting frequencies and adding each element to the heap.
4. Compute the x-sum for the initial window by taking the top x elements from the heap and summing their total contributions based on frequency.
5. Slide the window forward one element at a time:

a. Decrement the frequency of the element leaving the window. If its count reaches zero, remove it from the heap.

b. Increment the frequency of the element entering the window. Add it to the heap if it was not present.

c. Recompute the top x elements efficiently using the heap structure to obtain the new x-sum. 6. Append each x-sum to the result array. 7. Return the final array after processing all windows.

Why it works: The sliding window maintains the correct frequency counts as the window moves, and the heap or sorted structure guarantees we can efficiently access the top x most frequent elements with ties resolved by value. Each operation maintains the invariant that the heap correctly represents the current window's element frequencies, ensuring the sums are correct.

Python Solution

from collections import Counter
from sortedcontainers import SortedList
from typing import List

class Solution:
    def findXSum(self, nums: List[int], k: int, x: int) -> List[int]:
        n = len(nums)
        if x == k:
            return [sum(nums[i:i+k]) for i in range(n - k + 1)]
        
        result = []
        freq = Counter()
        sorted_list = SortedList()
        
        for i in range(k):
            val = nums[i]
            if freq[val] > 0:
                sorted_list.discard((-freq[val], -val))
            freq[val] += 1
            sorted_list.add((-freq[val], -val))
        
        for i in range(n - k + 1):
            sum_x = sum(-val * -count for count, val in sorted_list[:x])
            result.append(sum_x)
            
            if i + k < n:
                out_val = nums[i]
                sorted_list.discard((-freq[out_val], -out_val))
                freq[out_val] -= 1
                if freq[out_val] > 0:
                    sorted_list.add((-freq[out_val], -out_val))
                
                in_val = nums[i + k]
                if freq[in_val] > 0:
                    sorted_list.discard((-freq[in_val], -in_val))
                freq[in_val] += 1
                sorted_list.add((-freq[in_val], -in_val))
        
        return result

The code first checks for the trivial case where x equals k. It then uses a Counter to track element frequencies and a SortedList to maintain elements sorted by frequency descending and value descending. The heap-like operations on SortedList allow efficient addition and removal. For each window, we extract the top x elements from the SortedList and compute the x-sum, then slide the window by updating the frequency map and SortedList.

Go Solution

package main

import (
    "sort"
)

func findXSum(nums []int, k int, x int) []int64 {
    n := len(nums)
    result := make([]int64, 0, n-k+1)
    if x == k {
        for i := 0; i <= n-k; i++ {
            var sum int64
            for j := 0; j < k; j++ {
                sum += int64(nums[i+j])
            }
            result = append(result, sum)
        }
        return result
    }

    freq := map[int]int{}
    type pair struct{ val, count int }
    window := make([]pair, 0)

    for i := 0; i < k; i++ {
        freq[nums[i]]++
    }

    for i := 0; i <= n-k; i++ {
        window = window[:0]
        for num, count := range freq {
            window = append(window, pair{num, count})
        }
        sort.Slice(window, func(i, j int) bool {
            if window[i].count == window[j].count {
                return window[i].val > window[j].val
            }
            return window[i].count > window[j].count
        })
        var sum int64
        taken := 0
        for _, p := range window {
            for t := 0; t < p.count && taken < x; t++ {
                sum += int64(p.val)
                taken++
            }
            if taken >= x {
                break
            }
        }
        result = append(result, sum)

        if i+k < n {
            freq[nums[i]]--
            if freq[nums[i]] == 0 {
                delete(freq, nums[i])
            }
            freq[nums[i+k]]++
        }
    }

    return result
}

The Go solution mirrors the Python solution but uses slices and maps instead of SortedList. It constructs a sortable slice of pair structs for each window, sorts by frequency and value, and computes the top x contributions. Differences include manual slice management and type conversion to int64 to avoid overflow.

Worked Examples

Example 1: nums = [1,1,2,2,3,4,2,3], k = 6, x = 2

Initial window [1,1,2,2,3,4], frequencies: {1:2, 2:2, 3:1, 4:1}. Top 2: 2 and 1 → x-sum = 2+2+1+1 = 6

Next window [1,2,2,3,4,2], frequencies: {1:1,2:3,3:1,4:1}. Top 2: 2 and 4 → x-sum = 2+2+2+4 = 10

Next window [2,2,3,4,2,3], frequencies: {2:3,3:2,4:1}. Top 2: 2 and 3 → x-sum = 2+2+2+3+3 = 12

Example 2: nums = [3,8,7,8,7,5], k = 2, x = 2

All windows have k == x, so each sum equals the subarray sum: [3+8, 8+7, 7+8, 8+7, 7+5] = [11,15,15,15,12]

Complexity Analysis

| Measure | Complexity