SRM 601 Div 1 - Problem 250 WinterAndPresents http://community.topcoder.com/stat?c=problem_statement&pm=12860

WinterAndPresents.py

import itertools,sys

class WinterAndPresents:
    def getNumber(self, apple, orange):
        """This is an O(N) solution to SRM 601 Div 1 Problem 250.

        The derivation is as follows.
        Let :math:`a_i` and :math:`o_i` be the number of apples and oranges in bag :math:`i`, respectively.
        Let :math:`x_m` denote the largest possible value of :math:`x`.
        Given any :math:`x`, the number of different presents :math:`n(x)` can be written as:

        .. math::

            n(x) &= \sum_{i=1}^N \min(x, a_i) - \sum_{i=1}^N \max(o_i, x - o_i) + 1 \\\\
            &= \sum_i x \mathbb{1}[x \leq a_i] + \sum_i a_i \mathbb{1}[x > a_i] - \sum_i (x - o_i) \mathbb{1} [x_i > o_i] + 1 \\\\
            &= \sum_i x (1 - \mathbb{1}[x \leq a_i]) + \sum_i a_i \mathbb{1}[x > a_i] - \sum_i (x - o_i) \mathbb{1} [x_i > o_i] + 1 \\\\
            &= Nx - \sum_i (x - a_i) \mathbb{1} [x_i > a_i] - \sum_i (x - o_i) \mathbb{1} [x_i > o_i] + 1

        Therefore, the answer we are looking for is:

        .. math::

            \sum_{x=1}^{x_m} n(x) = N \sum_x x + \sum_x 1 - \sum_{i,x : x > a_i} (x-a_i) - \sum_{i,x : x > o_i} (x-o_i)

        It's easy to see that the first two terms takes O(1) to compute.
        By symmetry, the third and forth can be computed in a similar way.
        Take the third term for example:

        .. math::

            \sum_{i,x : x > a_i} (x-a_i)
            &= \sum_i \sum_{x: x > a_i} (x - a_i) \\\\
            &= \sum_i [\sum_{x: x > a_i} x - a_i \sum_{x: x > a_i} 1]

        Now given any :math:`i`, both terms in the brackets take O(1) to compute.
        Note that the key is to move the summation over :math:`i` before the summation over :math:`x`.

        """
        assert len(apple) == len(orange)
        N = len(apple)
        if N == 0:
            return 0

        max_x = sys.maxsize
        for a, o in itertools.izip(apple, orange):
            max_x = min(max_x, a + o)
        if max_x == 0:
            return 0

        number = max_x + N * (1 + max_x) * max_x / 2

        for i in xrange(N):
            def delta_given_fruit(fruit):
                return 0 if max_x <= fruit[i]\
                    else (fruit[i] + 1 + max_x) * (max_x - fruit[i]) / 2 - fruit[i] * (max_x - fruit[i])

            number -= delta_given_fruit(apple)
            number -= delta_given_fruit(orange)

        return number