Prison Cells After N Days

Intuition

Solution

class Solution:
    def key(self, cells):
        res = 0
        for i in cells:
            res = (res << 1) + i
        return res

    def value(self, key):
        res = [0] * 8
        for i in range(7, -1, -1):
            res[i] = key % 2
            key = key >> 1
        return res

    def child(self, key):
        cells = self.value(key)
        res = [0] * 8
        for i in range(1, 7):
            if cells[i - 1] == cells[i + 1]:
                res[i] = 1
        return self.key(res)

    def prisonAfterNDays(self, cells: List[int], N: int) -> List[int]:
        visited = [False] * 256
        prefix = []
        cycle = []

        index = self.key(cells)
        while not visited[index]:
            visited[index] = True
            prefix.append(index)
            index = self.child(index)

        cycle_start = index
        cycle.append(index)
        index = self.child(index)
        while index != cycle_start:
            cycle.append(index)
            index = self.child(index)

        for _ in cycle:
            prefix.pop()

        if N < len(prefix):
            return prefix[N]
        N -= len(prefix)
        return self.value(cycle[N % len(cycle)])