Problem statement

https://leetcode.com/problems/rle-iterator/

Solution

Basically, what we need to do is to iterate over our data and keep tuples (frequency, number). Let put them to stack. Then we need to do:

  1. If stack is not empty and n > self.stack[-1][0], it means, that we can remove last element from stack completely and update n.
  2. If stack is empty now, return -1 immediately.
  3. Extract last freq, num from stack and put updated frequency. Note, that freq - n >= 0, because in the opposite case we already managed it in while loop.
  4. Finally, return num.

Complexity

Time complexity is O(N + Q), where N is size of A and Q is number of next runs.

Code

class RLEIterator:
    def __init__(self, A):
        self.stack = [(A[2*i], A[2*i+1]) for i in range(len(A)//2 - 1, -1, -1)]

    def next(self, n):
        while self.stack and n > self.stack[-1][0]:
            n -= self.stack[-1][0]
            self.stack.pop()
            
        if not self.stack: return -1
        
        freq, num = self.stack.pop()     
        self.stack.append((freq-n, num))
        
        return num