# Locker puzzle in Python

(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.

• "I am not interested in the square numbers trick/optimization." — Why not? – Gareth Rees Apr 13 '15 at 13:42
• @GarethRees because I already know that in the end only the perfect squared doors remain open, but coding that would not be interesting. – Caridorc Apr 13 '15 at 20:02
• (To the downvoter) it is good practice to tell what to improve when downvoting. – Caridorc Apr 13 '15 at 20:04

## 1 Answer

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

• Sorry, the problem is now fixed – TheBlackCat Apr 13 '15 at 14:22
• Sorry, somehow misread the question as "how many are open". Fixed now. – TheBlackCat Apr 13 '15 at 14:29
• sigh, not doing very well here, sorry. Fixed. – TheBlackCat Apr 13 '15 at 14:32
• Okay, this version should work – TheBlackCat Apr 13 '15 at 14:35