lamps and streets problem

lamps.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: Yura A
"""


def rotate(sample):
    """
    поворачиваем образец по часовой стрелке.
    """
    result = []
    columns_number = len(sample[0])
    for column_x in xrange(columns_number):
        column_slice = ''.join(
            [row[column_x] for row in sample])[::-1]
        result.append(column_slice)
    return result


def light_to_right(sample, k):
    """
    заменяем елементы образца вправо от "2" на "-"
    на промежутке длиной k.
    """
    result = []
    for sample_row in sample:
        light_mode = None
        result_row = ''
        for cell in sample_row:
            result_cell = cell
            if '2' == cell:
                light_mode = 0
            elif '1' == cell:
                light_mode = None
            if light_mode is not None:
                if '0' == cell and light_mode <= k:
                    result_cell = '-'
                light_mode += 1
            result_row += result_cell
        result.append(result_row)
    return result


def the_dark(town, k):
    """
    вращаем город четыре раза.
    после каждого поворота обозначаем освещенные
    улицы вправо от фонаря.
    """
    for _ in xrange(4):
        town = rotate(town)
        town = light_to_right(town, k)
    # печатаем город
    for row in town:
        print row
    # выводим количество оставшихся нолей
    return sum([row.count('0') for row in town])


TOWN = [
    '0000000',
    '0000010',
    '0121000',
    '0000200',
    '0010000',
]


if __name__ == '__main__':
    print the_dark(TOWN, 2)