Python中的储物柜难题

Question

(游戏:储物柜之谜)学校有100个储物柜和100个学生。所有的储物柜在开学的第一天就关闭了。当学生们进入时，第一个学生，表示S1，打开每个储物柜。然后第二个学生，S2，从第二个储物柜开始，表示L2，关闭所有其他的储物柜。学生S3从第三个储物柜开始，并改变第三个储物柜(打开时关闭它，关闭时打开它)。学生S4从储物柜L4开始，每四个储物柜更换一次。学生S5从L5开始，每五个储物柜都换一个，等等，直到学生S100改变L100为止。所有的学生都经过大楼换了储物柜后，哪些储物柜是开着的？写一个程序来找到你的答案。

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

我对正方形的技巧/优化不感兴趣。

Answer 1

这类东西适合于numpy数组。您可以使用~arr反转布尔numpy数组，这是一个逐元素的not操作.您可以通过arr[start:stop:step] = ~arr[start:stop:step]对一个片段执行此操作，或者在您的情况下，您将使用arr[start::step] = ~arr[start::step]结束。这也是一个就地操作，从而避免不必要的副本.

因此，将代码直接转换为numpy将如下所示：

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

如果您能够避免列出列表，那么所述问题的一个更简单的版本是：

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