Moving Average from Data Stream

moving_average.py

"""
Description:
Given a stream of integers and a window size, calculate the moving average of all
integers in the sliding window.

Example:
m = MovingAverage(3)
m.next(1)   # 1
m.next(10)  # (1 + 10) / 2
m.next(3)   # (1 + 10 + 3) / 3
m.next(5)   # (10 + 3 + 5) / 3
"""

class MovingAverage:
    def __init__(self, size):
        """
        Initialize your data structure here.
        :type size: int
        """
        self.size = size
        self.sum = 0
        self.count = 0
        self.index = 0
        self.buff = [0] * size

    def next(self, val):
        """
        :type val: int
        :rtype: float
        """
        self.sum += val
        self.count += 1
        result = None
        if self.count > self.size:
            self.sum -= self.buff[self.index]
            result = self.sum / self.size
        else:
            result = self.sum / self.count
        self.buff[self.index] = val
        self.index += 1
        if self.index >= self.size:
            self.index = 0
        return result