MediumLeetCode 1352Design — prefix products with reset on zero

Product of the Last K Numbers (LeetCode 1352)

Problem 9 of 10 in Design a Data Structure — Practice Problems

The problem

Design a class that receives a stream of non-negative integers:

  • add(num): append num to the stream;
  • getProduct(k): return the product of the last k numbers added. There are always at least k numbers.
add(3), add(0), add(2), add(5), add(4)
getProduct(2) -> 20     5 × 4
getProduct(3) -> 40     2 × 5 × 4
getProduct(4) -> 0      0 × 2 × 5 × 4
add(8)
getProduct(2) -> 32     4 × 8

How to think about it

Start with brute force. Store every number, and for getProduct(k) multiply the last k. That is O(k) per query, which is too slow when there are many queries with large k.

The key insight: prefix products, like prefix sums. Keep a list pre where pre[i] is the product of the first i numbers (pre[0] = 1). Then the product of the last k numbers is pre[-1] // pre[-1 - k]: the product of everything, divided by the product of everything before the window. Both add and getProduct become O(1).

The catch: zero. Once a 0 enters, every later prefix product is 0, and you can no longer divide to recover a window. The fix: when a 0 arrives, reset pre to [1]. Now pre only covers the numbers after the most recent zero. If a query’s window fits inside that stretch (k < len(pre)), divide as normal. If the window reaches back past the reset (k >= len(pre)), it must contain that zero, so the answer is 0.

Which pattern is this? Prefix sums applied to products, inside a design question. The tell is “range product over the most recent elements of a stream”. Zeros are the twist that makes division unsafe, and the reset turns that problem into a length check.

Approach

  1. Start with pre = [1].
  2. add(num): if num is 0, reset pre = [1]. Otherwise append pre[-1] * num.
  3. getProduct(k): if k >= len(pre), the window includes a zero, so return 0. Otherwise return pre[-1] // pre[-1 - k].

Solution

class ProductOfNumbers:
    def __init__(self): self.pre = [1]
    def add(self, num):
        if num == 0: self.pre = [1]                      # zero: forget everything before it
        else: self.pre.append(self.pre[-1] * num)
    def getProduct(self, k):
        if k >= len(self.pre): return 0                  # window reaches the last zero
        return self.pre[-1] // self.pre[-1 - k]          # exact: pre[-1-k] divides pre[-1]
# add 3, 0, 2, 5, 4; getProduct(2) -> 20; getProduct(3) -> 40; getProduct(4) -> 0
# add 8; getProduct(2) -> 32

Dry run

The example sequence:

Operation Returns pre after
add(3) — [1, 3]
add(0) — [1] (reset)
add(2), add(5), add(4) — [1, 2, 10, 40]
getProduct(2) 40 // 2 = 20 [1, 2, 10, 40]
getProduct(3) 40 // 1 = 40 [1, 2, 10, 40]
getProduct(4) 0 (4 ≥ len 4) [1, 2, 10, 40]
add(8) — [1, 2, 10, 40, 320]
getProduct(2) 320 // 10 = 32 [1, 2, 10, 40, 320]

getProduct(4) needs four numbers but only three have arrived since the zero, so the window includes the 0.

Complexity

  • Time: O(1) per operation. One append or one division.
  • Space: O(n) for the prefix list, where n is the number of values added since the last zero.

Edge cases and pitfalls

  • Zeros are the whole difficulty. Without the reset, pre becomes all zeros and you divide by zero.
  • k == len(pre) also returns 0: you need k numbers but only len(pre) - 1 exist after the zero. Use >=, not >.
  • Integer division is exact here, because pre[-1 - k] is a factor of pre[-1]. Use // so the result stays an int.
  • Overflow in fixed-width languages. The problem only promises that each answer fits in 32 bits; the running prefix product can grow much larger. Python’s big integers hide this; in Java you need to discuss it.
  • Stream of ones makes pre long but constant; that is fine.