A wrapper class for calculating simple statistics over various sub-arrays of an array in amortized O(1) time

cumulativesum2D.py

class CumulativeSum2D:
  """A wrapper around 2D arrays to be able to quickly grab sums and averages of various sub-arrays of the array"""
  def __init__(self,data):
    self.cum_sum = self.__cumulify(data)
  
  def __cumulify(self,data):
    """ creates a 2D cumulative sum matrix"""
    cum_sum_array = []
    for j in xrange(len(data)):
      cum_sum_array.append([])
      for i in xrange(len(data[j])):
        cum_sum = 0
        for j_sub in xrange(0,j+1):
          for i_sub in xrange(0,i+1):
            cum_sum += data[j_sub][i_sub]
        cum_sum_array[j].append(cum_sum)
    return cum_sum_array
  
  def sum_over_range(self,x,y,width,height):
    """Given the coordinates of the top left point, and a width and length, calculate the sum over that area"""
    if x + width >= len(self.cum_sum[0]):
      raise IndexError("x coordinate is bigger than array")
    if y + height >= len(self.cum_sum):
      raise IndexError("y coordinate is bigger than array")
    total_sum = self.cum_sum[y+height][x+width]
    remove = 0
    if y != 0: remove += self.cum_sum[y-1][x+width]
    if x != 0: remove += self.cum_sum[y+height][x-1]
    if x != 0 and y != 0:
      remove -= self.cum_sum[y-1][x-1]
    return total_sum - remove
  
  def avg_over_range(self,x,y,width,height):
    """Given the coordinates of the top left point, and a width and length, calculate the average over that area"""
    return self.sum_over_range(x,y,width,height)/((width+1.0)*(height+1.0))

t = CumulativeSum2D([[1,2,3],
                     [4,5,6]])

assert(t.sum_over_range(0,0,1,1) == 12)
assert(t.sum_over_range(2,0,0,1) == 9)
assert(t.avg_over_range(2,0,0,1) == 4.5)
assert(t.sum_over_range(0,0,2,1) == 6*(6+1)/2.0)
assert(t.avg_over_range(0,0,2,1) == 6*(6+1)/(2*6.0))