Problem statement

https://binarysearch.com/problems/Prison-Cells/

Solution

Equal to Leetcode 0957. Prison Cells After N Days.

Complexity

It is O(64) for time and space.

Code

class Solution:
    def next_step(self, cells):
        res = [0] * 8
        for i in range(1,7):
            res[i] = int(cells[i-1] == cells[i+1])
        return res
    
    def solve(self, cells, N):
        found_dic = {}
        for i in range(N):
            cells_str = str(cells)
            if cells_str in found_dic:
                loop_len = i - found_dic[cells_str]
                return self.solve(cells, (N - i) % loop_len)
            else:
                found_dic[cells_str] = i 
                cells = self.next_step(cells) 
                
        return cells