LESSON 313 OF 414 · SLIDING WINDOW

Longest Subarray with at Most K Distinct Integers in Scala

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

SCALA · SLIDING WINDOW

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

Problem overview

The method longestSubarray receives an Array[Int] called nums and an integer K. It returns the length of the longest contiguous piece of nums (a subarray) that contains at most K different values. Repeats of a value inside the piece don't count twice: [1, 2, 1, 2] has four elements but only two distinct values.

For nums = [1, 2, 1, 2, 3] and K = 2, the piece [1, 2, 1, 2] has two distinct values, and adding the 3 would make three. No longer valid piece exists, so the answer is 4.

Approach: a window that grows and shrinks

The brute force looks at every start and end position and counts the distinct values in each piece, which takes at least O(n²) time. The sliding window does much better because of one property: if a window has too many distinct values, removing elements from its left end is the only way to fix it, and a window that is valid stays valid when you remove elements from it.

So the code keeps a window nums[first..last]. For every position last it extends the window one element to the right, then moves first to the right until the window has at most K distinct values again. The best window length seen along the way is the answer. Both ends only move forward, so each index enters the window once and leaves it at most once, which makes the whole scan O(n).

To know how many distinct values the window holds, the code keeps a map from each value to how many times it occurs in the window. The map's size is the number of distinct values, provided that a value is deleted from the map as soon as its count reaches zero.

Here is the example with K = 2. For last from 0 to 3 the window grows to [1, 2, 1, 2] and the map holds 1 -> 2, 2 -> 2, so the answer is 4. At last = 4 the value 3 enters and the map has three keys. The loop removes nums(0) (the count of 1 drops to 1, still three keys), then nums(1) (the count of 2 drops to 1, still three keys), then nums(2) (the count of 1 reaches 0, so the key disappears and two remain). Now first = 3 and the window is [2, 3] with length 2, which doesn't beat 4.

Step by step

Create the frequency map as a val holding an empty mutable.Map[Int, Int]. The val only fixes which map the name refers to; the map's contents can still change.
Declare answer and first as vars set to 0. answer is the best length so far and first is where the window starts.
Loop over last in nums.indices. Each index is tried once as the right end of the window.
Call insert(freq, nums(last)) to add the new element. Arrays are indexed with parentheses in Scala.
Shrink while the window is invalid: a while loop runs as long as freq.size > K. A single removal can leave the count unchanged (when the removed value still occurs inside the window), so this must be a loop, not an if.
Remove the leftmost element with remove(freq, nums(first)) and move first one step right with first += 1. Scala has no ++.
Record the length. After the loop the window is valid, so answer becomes math.max(answer, last - first + 1).
Return answer after the loop.
Write remove with the parameters freq: mutable.Map[Int, Int] and value: Int, returning Unit. It lowers the count by one.
Delete the key at zero with freq.remove(value). Without this, freq.size would keep counting values that have left the window.
Write insert with the same parameters. It stores the old count plus one, reading the old count with freq.getOrElse(value, 0) because a new value has no key yet.

The solution in Scala

class Solution {
    def insert(freq: mutable.Map[Int, Int], value: Int): Unit = {
        freq(value) = freq.getOrElse(value, 0) + 1
    }

    def remove(freq: mutable.Map[Int, Int], value: Int): Unit = {
        freq(value) = freq(value) - 1
        if (freq(value) == 0) {
            freq.remove(value)
        }
    }

    def longestSubarray(nums: Array[Int], K: Int): Int = {
        val freq = mutable.Map[Int, Int]()
        var answer = 0
        var first = 0
        for (last <- nums.indices) {
            insert(freq, nums(last))
            while (freq.size > K) {
                remove(freq, nums(first))
                first += 1
            }
            answer = math.max(answer, last - first + 1)
        }
        return answer
    }
}

The Scala details worth a closer look:

Importing mutable. The code writes mutable.Map, which lives in the package scala.collection.mutable. In a standalone file you need import scala.collection.mutable at the top. Plain Map without the prefix is the immutable one, which has no update or remove.
freq(value) = ... is a method call. Assigning to freq(value) calls freq.update(value, ...), which adds the key or replaces its value.
Reading a missing key throws. freq(7) on a map without the key throws NoSuchElementException: key not found: 7. getOrElse(value, 0) returns the default, and get(value) returns an Option. remove never reads a missing key, because it is only called for values inside the window.
freq.remove(value) deletes the key and returns the removed count wrapped in an Option. Here that result is simply discarded.
return at the end. The last expression of a block is its value, so answer alone would do. The explicit return answer also works.

Complexity

Time: O(n). The outer loop runs n times. The inner while loop moves first forward, and first never moves back or passes n, so across the whole run it executes at most n times. Each map operation takes constant time on average.

Space: O(n) in the worst case, more precisely O(min(n, K + 1)). The map holds one key per distinct value in the window, and it never has more than K + 1 keys because the window is shrunk as soon as it reaches that many.

Edge cases and common mistakes

Empty array. The loop never runs and the method returns 0.
K = 0. Every element makes the map non-empty, so the window is emptied again at once and the answer is 0.
K at least the number of distinct values. Nothing is ever removed and the answer is nums.length. For [1, 2, 3, 4] and K = 10 it is 4.
All values equal. [4, 4, 4, 4] with K = 1 returns 4.
Leaving zero counts in the map. If remove only decrements, a value with count 0 still counts toward freq.size, the window shrinks too much and the answer comes out too small.
Using if instead of while. In the worked example, three removals are needed before the window is valid again. One removal leaves it invalid.
Counting the length wrongly. The window nums[first..last] includes both ends, so its length is last - first + 1.
freq(value) += 1 for a new value. It reads the key first, so it throws for a value not yet in the map. That is why insert uses getOrElse.

Testing

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

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

Start learning free