SRM493 - div2 - level one

srm493_div2_level1.py

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

"""SRM493 - div2 - level one

Little Romeo likes cosmic amoebas a lot. Recently he received one as a gift from 
his mother. He decided to place his amoeba on a rectangular table. The table is a 
grid of square 1x1 cells, and each cell is occupied by either matter or antimatter. 
The amoeba is a rectangle of size 1xK. Romeo can place it on the table in any 
orientation as long as every cell of the table is either completely covered by part 
of the amoeba or completely uncovered, and no part of the amoeba lies outside of 
the table. It is a well-known fact that cosmic amoebas cannot lie on top of matter, 
so every cell of the table covered by the amoeba must only contain antimatter.

You are given a String[] table, where the j-th character of the i-th element is 
'A' if the cell in row i, column j of the table contains antimatter or 'M' if it 
contains matter. Return the number of different ways that Romeo can place the cosmic 
amoeba on the table. Two ways are considered different if and only if there is a 
table cell that is covered in one but not the other.

Definition:
    Class: AmoebaDivTwo
    Method: count
    Parameters: String[], int
    Returns: int
    Method signature: int count(String[] table, int K)
"""

class AmoebaDivTwo(object):
    """Amoeba Div Two

    Constraints:
        - table will contain between 1 and 50 elements, inclusive.
        - Each element of table will contain between 1 and 50 characters, inclusive.
        - All elements of table will have the same length.
        - Each character of each element of table will be either 'A' or 'M'.
        - K will be between 1 and 50, inclusive.

    >>> amoeba = AmoebaDivTwo()
    >>> print amoeba.count(["MA", ], 2)
    0
    >>> print amoeba.count(["AAA", "AMA", "AAA", ], 3)
    4
    >>> print amoeba.count(["AA", "AA", "AA", ], 2)
    7
    >>> print amoeba.count(["MMM", "MMM", "MMM", ], 1)
    0
    >>> print amoeba.count(["AAM", "MAM", "AAA", ], 1)
    6
    >>> print amoeba.count(["AAA", "AAM", "AAA", ], 2)
    9
    """

    def __init__(self):
        pass

    def count(self, table=[], K=1):
        c = 0
        for y in range(len(table)):
            for x in range(len(table[0]) - K + 1):
                if table[y][x:x+K] == 'A'*K:
                    c += 1
        if K <= 1:
            return c

        table = [''.join([table[y][x] for y in range(len(table))]) for x in range(len(table[0]))]
        for y in range(len(table)):
            for x in range(len(table[0]) - K + 1):
                if table[y][x:x+K] == 'A'*K:
                    c += 1
        return c

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