Leetcode 994 Solution

This article provides solution to leetcode question 994 (prison-cells-after-n-days)

https://leetcode.com/problems/prison-cells-after-n-days

Solution

class Solution:
    def prisonAfterNDays(self, cells: List[int], N: int) -> List[int]:
        def nextday(cells):
            new_cells = [0] * len(cells)
            for i in range(1, len(new_cells) - 1):
                new_cells[i] = 1 if cells[i - 1] == cells[i + 1] else 0
            return new_cells

        seen = {}
        while N:
            if tuple(cells) in seen:
                N %= seen[tuple(cells)] - N
            else:
                seen[tuple(cells)] = N

            if N >= 1:
                cells = nextday(cells)
                N -= 1
        return cells