LESSON 283 OF 338 · SLIDING WINDOW

Longest Subarray With At Most K Distinct Integers in Go: Sliding Window

Given an array of integers, find the longest subarray that contains at most K distinct integers and return its length.

Example

Input: nums = [1, 2, 1, 2, 3], k = 2
Output: 4
Explanation:the subarray nums[0...3] contains 2 distinct values: [1, 2] and is the longest subarray

Note:

Your algorithm should run in O(n^2) time and use O(n) extra space.

STEP 1 OF 14

Guided

Ctrl + Enter to check

GO · SLIDING WINDOW

Longest Subarray With At Most K Distinct Integers in Go, explained

Problem overview

The function longestSubarray receives a slice nums of integers and a number k. It returns the length of the longest contiguous subarray that contains at most k different values. Repeats of the same value count once: [1, 2, 1, 2] has two distinct values although it has four elements.

In the example, nums = []int{1, 2, 1, 2, 3} and k = 2. The first four elements contain only the values 1 and 2, so a window of length 4 is allowed. Adding the 3 would make three distinct values, and the window [2, 3] at the end is shorter, so the answer is 4.

The statement allows O(n²) time and asks for O(n) extra space. The approach below needs only O(n) time.

Approach: grow the window, shrink it when it has too many kinds

The simple solution is to fix each start and extend to the right, tracking the distinct values in a set, until a value number k + 1 shows up. That scans up to n elements for each of n starts, which is the O(n²) budget.

The faster idea changes who moves. Instead of fixing the start, fix the end of the window and let the start follow:

Visit each index last in turn and add nums[last] to the window.
If the window now has more than k distinct values, remove elements from the front (index first) until it has at most k again.
The window is now the longest valid one ending at last, so compare its length with the best so far.

This works because removing elements from the front can only reduce the number of distinct values, and a valid window stays valid when it loses elements. The start never needs to move backwards, since a window that was too wide for an earlier end is also too wide for a later one.

The window needs more than a sum. To know when a value really leaves, count the occurrences of each value in a map frequency. Adding a value increments its count. Removing a value decrements it, and only when the count reaches zero is the key deleted. Then len(frequency) is exactly the number of distinct values in the window.

On the example with k = 2:

After indices 0 to 3, frequency holds two keys, 1 and 2, each with count 2, and the window length is 4.
Index 4 adds the value 3, which makes three keys. Removing nums[0] leaves counts 1:1, 2:2, 3:1, still three keys. Removing nums[1] still leaves three, because the value 2 occurs again. Removing nums[2] deletes the key 1, leaving two keys, and first is now 3. The window is [2, 3] with length 2.
The best length seen is 4.

Step by step

Create the map frequency with make. A map that is only declared is nil, and writing to it panics.
Declare answer and first as 0: the best length so far and the start of the window.
Loop over last from 0 to len(nums) - 1, the end of the window.
Insert nums[last] with the helper insert.
Shrink while there are too many kinds. Go has no while, so use a for with only the condition len(frequency) > k.
Remove nums[first] with the helper remove, then move first forward. Each index leaves the window at most once, so this inner loop is cheap in total.
Update answer with the length last - first + 1, using the built-in max.
Return answer after the loop.
Write remove. It takes the map and a value and returns nothing. It decrements the value's count, and when the count hits 0, deletes the key.
Write insert. It increments the value's count. A missing key reads as 0, so the first ++ creates the entry with count 1.

The solution in Go

func longestSubarray(nums []int, k int) int {
    frequency := make(map[int]int)
    answer, first := 0, 0
    for last := 0; last < len(nums); last++ {
        insert(frequency, nums[last])
        for len(frequency) > k {
            remove(frequency, nums[first])
            first++
        }
        answer = max(answer, last-first+1)
    }
    return answer
}

func remove(frequency map[int]int, value int) {
    frequency[value]--
    if frequency[value] == 0 {
        delete(frequency, value)
    }
}

func insert(frequency map[int]int, value int) {
    frequency[value]++
}

The Go details worth knowing:

Maps are reference types. insert and remove receive the map by value, but that value is a handle to the same underlying table, so changes are visible to the caller without a pointer.
Missing keys read as zero. frequency[value]++ on an absent key reads 0, adds 1 and stores 1. That is why insert needs no if.
Delete is required. A key left at count 0 still counts in len(frequency), so the window would look as if it held values it no longer contains. delete on a missing key is harmless.
Helpers can come after their callers. Go lets a function use functions declared anywhere in the same package, so remove and insert can sit below longestSubarray.
Built-in max. It has been part of the language since Go 1.21, so no helper is needed for two integers.
Map iteration order is random, but this solution never iterates the map. It only uses len, so the order doesn't matter.

Complexity

Time: O(n). last visits each index once. first only moves forward and never passes last, so across the whole run the inner loop removes each element at most once. Map operations take constant time on average, so the total is linear.

Space: O(min(k, n)). After shrinking, the map holds at most k keys, and during an insert it may briefly hold k + 1. It can never hold more keys than there are elements, so the memory stays within the O(n) the statement allows.

Edge cases and common mistakes

Empty slice. The loop never runs and the function returns 0.
k is 0. Every insert makes the map too big, so the window shrinks to empty each time, and the answer is 0. For nums = []int{1, 2, 1, 2, 3} and k = 0 the function returns 0.
k at least the number of distinct values. The window never shrinks and the result is len(nums): for the example with k = 3 or k = 10 it returns 5.
All values equal. nums = []int{4, 4, 4, 4} and k = 1 returns 4, because the map has a single key whose count grows.
Forgetting to delete a zero count. The map then keeps a key for every value ever seen, len(frequency) never drops, and the shrinking loop runs first past the end of the slice. For the example with k = 2 the program stops with an index out of range panic.
Deleting the key on every decrement. A value that occurs several times in the window would be forgotten after its first removal. For nums = []int{1, 1, 1, 2, 3, 2} and k = 2 that version returns 5 instead of 4, because it lets the window hold three kinds.
Using an if instead of a for to shrink. One removal may not be enough: in the example, three removals are needed at index 4.
Reading k as the window size. k limits the number of different values, not the number of elements.

Testing

nums = []int{1, 2, 1, 2, 3}, k = 2 gives 4.
nums = []int{1, 2, 1, 2, 3}, k = 1 gives 1.
nums = []int{1, 2, 3, 4}, k = 2 gives 2.
nums = []int{4, 4, 4, 4}, k = 1 gives 4.
nums = []int{}, k = 2 gives 0.

Keep going with GoAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.

Start learning free