SRM485 - div2 - level two

srm485_div2_level2.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-

"""SRM485 - div2 - level two

A sequence of integers a[0], a[1], ..., a[N-1], where N >= 3, 
is called an arithmetic progression if the difference between 
any two successive terms in the sequence is a constant. 
More precisely, it is an arithmetic progression 
if a[i] - a[i-1] = a[i+1] - a[i] for every integer i such 
that 0 < i < N-1.

Sasha and Pasha are students sharing the same room. Once, when 
Pasha had left the room, Sasha was in a good mood. So he took 
a piece of chalk and wrote an arithmetic progression on the blackboard. 
The progression consisted of at least 4 elements, all of 
which were positive integers. Then Sasha went to class, 
and Pasha came back.

Pasha is a very nice person, but there's one problem with him
-- he's frightened of even numbers! So, when he returned, 
he decided to make all the numbers on the board odd. He did 
this by repeatedly finding an arbitrary even number X, erasing it, 
and writing X/2 in its place. He continued to perform 
this step until no even numbers remained.

Your task is to help Sasha restore the initial progression. 
You will be given a int[] seq, where the i-th element is 
the i-th number in the sequence written on the blackboard 
after Pasha's actions. Return an int[] whose i-th element is 
the i-th number in a sequence that Sasha could have originally 
written on the blackboard. The constraints will ensure 
that at least one such sequence exists. If there are several 
such sequences, choose the one among them whose int[] 
representation is lexicographically smallest.

Definition:
    Class: AfraidOfEven
    Method: restoreProgression
    Parameters: int[]
    Returns: int[]
    Method signature: int[] restoreProgression(int[] seq)

Notes:
    - The lexicographically smaller of two different int[]s 
      containing the same number of elements is the one 
      with a smaller number at the first position where they differ.
    - It is guaranteed that all elements of the resulting int[] 
      will fit into a 32-bit signed integer data type.
"""

class AfraidOfEven(object):
    """Afraid Of Even

    Constraints:
        - seq will contain between 4 and 50 elements, inclusive.
        - Each element of seq will be between 1 and 1000, inclusive.
        - Each element of seq will be odd.
        - There will exist at least one arithmetic progression from 
          which Pasha would produce exactly the sequence described by seq.

    >>> afraid = AfraidOfEven()
    >>> print afraid.restoreProgression([1, 1, 3, 1, 5, ])
    [1, 2, 3, 4, 5]
    >>> print afraid.restoreProgression([9, 7, 5, 3, 1, ])
    [9, 7, 5, 3, 1]
    >>> print afraid.restoreProgression([999, 999, 999, 999, ])
    [999, 999, 999, 999]
    >>> print afraid.restoreProgression([7, 47, 5, 113, 73, 179, 53, ])
    [14, 47, 80, 113, 146, 179, 212]
    >>> print afraid.restoreProgression([749, 999, 125, 1, ])
    [1498, 999, 500, 1]
    """

    def __init__(self):
        pass

    def restoreProgression(self, seq):
        def check(l, seq):
            for i in range(len(seq)):
                while l[i] > 0 and l[i] % 2 == 0:
                    l[i] /= 2
                if l[i] != seq[i]:
                    return False
            return True

        foo = []
        bar = []
        for i in range(len(seq)):
            foo.append(seq[0]+(seq[2]-seq[0])/2*i)
            bar.append(seq[1]+(seq[3]-seq[1])/2*(i-1))
        if check(foo[:], seq) and check(bar[:], seq):
            return min(foo, bar)
        if check(foo[:], seq):
            return foo
        return bar

if __name__ == '__main__':
    import doctest
    doctest.testmod()