Locker puzzle in Python

Question

(Game: locker puzzle) A school has 100 lockers and 100 students. All lockers are closed on the first day of school. As the students enter, the first student, denoted S1, opens every locker. Then the second student, S2, begins with the second locker, denoted L2, and closes every other locker. Student S3 begins with the third locker and changes every third locker (closes it if it was open, and opens it if it was closed). Student S4 begins with locker L4 and changes every fourth locker. Student S5 starts with L5 and changes every fifth locker, and so on, until student S100 changes L100.

After all the students have passed through the building and changed the lockers, which lockers are open? Write a program to find your answer.

import doctest

def invert_multiples(bools, n):
    """
    Inverts all the items at index n*a

    >>> invert_multiples([True, True, True, True], 2)
    [False, True, False, True]
    """
    return [not bool_ if index % n == 0 else bool_
                      for index, bool_ in enumerate(bools)
                      ]

def invert_all_multiples(bools, nums):
    """
    Subsequentially inverts all the items that are
    multiples of each num. 

    >>> doors = invert_all_multiples([True]*100, range(1,101))
    >>> [i[0] for i in enumerate(doors) if not i[1]]
    [1, 4, 9, 16, 25, 36, 49, 64, 81]
    """
    for num in nums:
        bools = invert_multiples(bools, num)
    return bools

if __name__ == "__main__":
    doctest.testmod()

I am not interested in the square numbers trick/optimization.

Answer 1

This sort of thing is suited to numpy arrays. You can invert a boolean numpy array using ~arr, which is an element-by-element not operation. You can do this to a slice by arr[start:stop:step] = ~arr[start:stop:step], or in your case where you are going to the end arr[start::step] = ~arr[start::step]. It is also an in-place operation, thus avoiding unnecessary copies.

So a direct translation of your code to numpy would be like so:

import numpy as np

def invert_multiples(bools, n):
    """
    Inverts all the items at index n*a

    >>> invert_multiples([True, True, True, True], 2)
    [False, True, False, True]
    """
    bools[num::num+1] = ~bools[num::num+1]

def invert_all_multiples(bools, nums):
    """
    Subsequentially inverts all the items that are
    multiples of each num. 

    >>> doors = invert_all_multiples([True]*100, range(1,101))
    >>> [i[0] for i in enumerate(doors) if not i[1]]
    [1, 4, 9, 16, 25, 36, 49, 64, 81]
    """
    bools = np.array(bools)
    for num in nums:
        invert_multiples(bools, num)
    return bools.tolist()

If you are able to avoid lists, a simpler version of the stated problem is:

nmax = 100
lockers = np.zeros(nmax, dtype='bool')
for n in range(nmax):
    lockers[n::n+1] = ~lockers[n::n+1]
res = np.flatnonzero(lockers)  # add one to get one-based indexing