LeetCode 900: RLE Iterator

Problem Restatement

We need to design an iterator over a run-length encoded sequence.

The encoded array stores pairs:

[count, value]

At every even index i, encoding[i] is the number of times encoding[i + 1] appears in the decoded sequence.

For example:

encoding = [3, 8, 0, 9, 2, 5]

represents:

[8, 8, 8, 5, 5]

because there are three 8s, zero 9s, and two 5s.

The iterator supports:

next(n)

This consumes the next n elements and returns the last consumed element.

If there are fewer than n elements left, consume what remains and return -1.

LeetCode defines next(n) this way, with n >= 1, and guarantees at most 1000 calls per test case.

Input and Output

Class shape:

class RLEIterator:

    def __init__(self, encoding: list[int]):
        ...

    def next(self, n: int) -> int:
        ...

Examples

Example:

Input:
["RLEIterator","next","next","next","next"]
[[[3,8,0,9,2,5]],[2],[1],[1],[2]]

Output:
[null,8,8,5,-1]

The encoded sequence is:

[8, 8, 8, 5, 5]

Operations:

The last call returns -1 because it asks for two elements, but only one element remains.

First Thought: Decode the Whole Sequence

A simple idea is to expand the encoded array into a normal list.

For:

[3, 8, 0, 9, 2, 5]

we would build:

[8, 8, 8, 5, 5]

Then next(n) just moves a pointer forward.

This is easy, but it can use too much memory because counts can be very large. The encoded array may be short while the decoded sequence is huge.

Key Insight

We only need to know the current run.

A run is one pair:

encoding[index]     = remaining count
encoding[index + 1] = value

When next(n) is called, we consume from the current run.

If the current run has fewer than n elements, we skip the whole run and subtract its count from n.

If the current run has at least n elements, we subtract n from its count and return its value.

So the iterator can be represented by:

Algorithm

Constructor:

1. Store encoding.
2. Set index = 0.

For next(n):

1. While index is valid and the current count is smaller than n:
   1. Subtract the current count from n.
   2. Move index to the next run by adding 2.
2. If index has reached the end, return -1.
3. Otherwise, the current run has enough elements.
4. Subtract n from the current count.
5. Return the current value.

Walkthrough

Use:

encoding = [3, 8, 0, 9, 2, 5]

Start:

index = 0
current run = 3 copies of 8

Call next(2):

Encoding is now logically:

[1, 8, 0, 9, 2, 5]

Call next(1):

Call next(1):

Call next(2):

Correctness

The iterator maintains the invariant that index points to the current run, and encoding[index] stores the number of unconsumed elements left in that run.

When next(n) is called, if the current run has fewer than n remaining elements, then all elements in that run must be consumed before the nth requested element can be reached. The algorithm subtracts that count from n and moves to the next run, preserving the meaning of n as the number of elements still needed.

If the algorithm reaches the end of the encoded array, then all remaining runs have been consumed and fewer than the requested n elements existed. Returning -1 is correct.

Otherwise, the current run contains at least n elements. The nth consumed element lies inside this run, so its value is encoding[index + 1]. The algorithm subtracts n from the remaining count and returns that value.

Thus, every call consumes exactly the requested number of elements when possible, returns the last consumed value, and returns -1 when not enough elements remain.

Complexity

Let:

m = len(encoding)

Space complexity is O(1) extra space, not counting the stored input array.

Implementation

class RLEIterator:

    def __init__(self, encoding: list[int]):
        self.encoding = encoding
        self.index = 0

    def next(self, n: int) -> int:
        while self.index < len(self.encoding) and self.encoding[self.index] < n:
            n -= self.encoding[self.index]
            self.index += 2

        if self.index == len(self.encoding):
            return -1

        self.encoding[self.index] -= n
        return self.encoding[self.index + 1]

Code Explanation

The constructor stores the encoded representation:

self.encoding = encoding
self.index = 0

The pointer index always points to a count position.

In next(n), we skip runs that do not have enough remaining elements:

while self.index < len(self.encoding) and self.encoding[self.index] < n:
    n -= self.encoding[self.index]
    self.index += 2

If there are no runs left, the request cannot be satisfied:

if self.index == len(self.encoding):
    return -1

Otherwise, the current run contains the last consumed element:

self.encoding[self.index] -= n
return self.encoding[self.index + 1]

We reduce the remaining count and return the run value.

Testing

def run_tests():
    it = RLEIterator([3, 8, 0, 9, 2, 5])

    assert it.next(2) == 8
    assert it.next(1) == 8
    assert it.next(1) == 5
    assert it.next(2) == -1

    it = RLEIterator([1, 10])
    assert it.next(1) == 10
    assert it.next(1) == -1

    it = RLEIterator([0, 1, 0, 2, 3, 3])
    assert it.next(2) == 3
    assert it.next(1) == 3
    assert it.next(1) == -1

    it = RLEIterator([5, 7])
    assert it.next(3) == 7
    assert it.next(2) == 7
    assert it.next(1) == -1

    print("all tests passed")

run_tests()

Test meaning: