Coding Dojo in Python

dojo.py

# -*- coding: utf-8 -*-
'''
http://dojopuzzles.com/problemas/exibe/lampadas-no-corredor/
https://pypi.python.org/pypi/dose
05/03/15

Um homem chamado José é o responsável por ligar e desligar as luzes de um corredor. Cada lâmpada tem seu próprio interruptor que liga e a desliga. Inicialmente todas as lâmpadas estão desligadas.
José faz uma coisa peculiar: se existem n lâmpadas no corredor, ele caminha até o fim do corredor e volta n vezes. Na iésima caminhada, ele aperta apenas os interruptores aos quais sua posição é divisível por i. Ele não aperta nenhum interruptor na volta à sua posição inicial, apenas na ida. A iésima caminhada é definida como ir ao fim do corredor e voltar.
Determine qual é o estado final de cada lâmpada. Está ligada ou desligada?
Exemplo:
Entrada: 3
Saída: [on, off, off]

Linha para rodar os testes
$ sudo pip install python-pytest
py.test dojo.py

True * 11
True * n
'''

'''
def corredor(n):
    lampadas = [True] * n
    for i in range(2, n + 1):
        for j in range(n):
            if (j + 1) % i == 0:
                lampadas[j] = not lampadas[j]
    return lampadas
'''


def corredor(n):
    lampadas = [True] * n
    for i in range(2, n + 1):
        for j in range(i - 1, n, i):
            lampadas[j] = not lampadas[j]
    return lampadas


def test_6_lampadas():
    assert corredor(6) == [True, False, False, True, False, False]


def test_5_lampadas():
    assert corredor(5) == [True, False, False, True, False]


def test_4_lampadas():
    assert corredor(4) == [True, False, False, True]


def test_3_lampadas():
    assert corredor(3) == [True, False, False]


def test_2_lampadas():
    assert corredor(2) == [True, False]


def test_1_lampada():
    assert corredor(1) == [True]


def test_nenhuma_lampada():
    entrada = 0
    obtido = corredor(entrada)
    assert obtido == []