LESSON 314 OF 415 · SLIDING WINDOW

Longest Subarray with at most K Distinct Integers in Swift

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 at most O(n²) time and use O(n) extra space.

STEP 1 OF 13

Guided

Ctrl + Enter to check

SWIFT · SLIDING WINDOW

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

Problem overview

The method longestSubarray receives an array of integers, nums, 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, so [1, 2, 1, 2] has two distinct values even though it has four elements.

In the example, nums = [1, 2, 1, 2, 3] and K = 2. The first four elements, [1, 2, 1, 2], hold only the values 1 and 2, so the answer is 4. Adding the 3 would make three distinct values, which is too many.

Approach: a window that grows and shrinks

The slow way is to check every subarray and count its distinct values. That repeats a lot of work, because neighbouring subarrays share almost all of their elements.

A sliding window avoids this. Keep a window nums[first...last] and move its right end last across the array one index at a time. Each new element may push the number of distinct values above K. When that happens, move first to the right, dropping elements from the left, until the window is valid again. After that, the window ending at last is the longest valid one that ends there, so its length is a candidate for the answer.

The window needs a way to tell how many distinct values it holds. A dictionary from value to count does it: freq[v] is the number of times v occurs in the window, and freq.count is the number of keys, which is the number of distinct values. For this to work, a value whose count drops to 0 must be removed from the dictionary, not left there with a zero.

Step by step

Create the counts. freq is an empty [Int: Int]. It is a var because the helpers change it.
Create the answer and the left edge. Both answer and first start at 0, declared together on one var line.
Loop over the right edge. for last in 0..<nums.count tries every index as the end of a subarray.
Add the new element. Call insert(&freq, nums[last]). The window now covers first...last.
Shrink while the window is invalid. A while loop runs as long as freq.count > K.
Drop the leftmost element. Inside the loop, call remove(&freq, nums[first]) and then do first += 1. Swift has no ++ operator.
Record the length. Once the loop ends, the window is valid, and its length is last - first + 1. Keep the larger of that and answer with max.
Return answer after the loop.
Write remove. It takes the dictionary as an inout [Int: Int] and a value, both with the _ label so the calls don't name the arguments. It subtracts one from the value's count, and if the count reaches 0 it deletes the key.
Write insert. It has the same parameters and adds one to the value's count, starting from 0 when the key is missing.

The solution in Swift

class Solution {
    func insert(_ freq: inout [Int: Int], _ val: Int) {
        freq[val, default: 0] += 1
    }

    func remove(_ freq: inout [Int: Int], _ val: Int) {
        freq[val]! -= 1
        if freq[val] == 0 {
            freq.removeValue(forKey: val)
        }
    }

    func longestSubarray(_ nums: [Int], _ K: Int) -> Int {
        var freq = [Int: Int]()
        var answer = 0, first = 0
        for last in 0..<nums.count {
            insert(&freq, nums[last])
            while freq.count > K {
                remove(&freq, nums[first])
                first += 1
            }
            answer = max(answer, last - first + 1)
        }
        return answer
    }
}

The Swift details worth a closer look:

inout and &. A dictionary is a value type, so passing it normally gives the function a copy. Marking the parameter inout and writing &freq at the call lets the helper change the caller's dictionary.
default: 0. Reading freq[val] gives an optional, which is nil for a missing key. The form freq[val, default: 0] += 1 reads 0 for a missing key and stores the new count, so insert needs no if.
The ! in remove. freq[val]! -= 1 force-unwraps the optional. It is safe here because only values currently in the window are removed. On a missing key it would crash the program.
Comparing an optional with 0. freq[val] == 0 compares an Int? with a literal, which Swift allows. It is false for a missing key and true only when the count is 0.
Ignored return value. removeValue(forKey:) returns the old value, but Swift doesn't warn about ignoring it because the method is marked as discardable.
Why delete the zero. If a count of 0 stayed in the dictionary, freq.count would include a value that is no longer in the window, and the window would shrink too far.

Complexity

Time: O(n). The index last visits each position once. The index first only moves right and never passes last, so across the whole run the inner loop does at most n removals. Each dictionary operation takes constant time on average.

Space: O(n) in the worst case. The dictionary holds at most K + 1 keys, because the window shrinks as soon as it exceeds K, and it can never hold more keys than the array has elements.

Edge cases and common mistakes

K = 0. No value is allowed, so the window is emptied after every step and the method returns 0. For [1, 2, 3] with K = 0 it returns 0.
Empty array. The loop never runs, so the method returns 0.
Large K. When K is at least the number of distinct values, the window never shrinks and the answer is nums.count. For [1, 2, 3, 4] with K = 10 the result is 4.
One value repeated. [4, 4, 4] with K = 1 returns 3.
Leaving zero counts behind. Decrementing without deleting the key at 0 makes freq.count too large.
Forgetting &. Calling insert(freq, nums[last]) is a compile error, because an inout argument must be passed with &, and freq must be a var.
Off-by-one length. Both ends are inside the window, so the length is last - first + 1, not last - first.

Testing

nums = [1, 2, 1, 2, 3], K = 2 gives 4.
nums = [1, 2, 3], K = 0 gives 0.
nums = [], K = 3 gives 0.
nums = [4, 4, 4], K = 1 gives 3.
nums = [1, 2, 3, 4], K = 10 gives 4.

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

Start learning free