题目描述

Time O(1)
Space O(1)

class Solution:
    def prisonAfterNDays(self, cells: List[int], N: int) -> List[int]:
        #0->0: 0
        #0->1: 2
        #1->0: 3
        #1->1: 1

        def change_per_day():
            if cells[0] == 1:
                cells[0] = 3
            if cells[-1] == 1:
                cells[-1] = 3

            for j in range(1, len(cells) - 1):
                if ((cells[j-1] == 0 or cells[j-1] == 2) and  (cells[j+1] == 0 or cells[j+1] == 2)) or ((cells[j-1] == 1 or cells[j-1] == 3) and  (cells[j+1] == 1 or cells[j+1] == 3)):
                    if cells[j] == 0:
                        cells[j] = 2
                else:
                    if cells[j] == 1:
                        cells[j] = 3

            for i in range(len(cells)):
                if cells[i] == 3:
                    cells[i] = 0
                elif cells[i] == 2:
                    cells[i] = 1

        change_per_day() #day 1

        mod = N % 14

        if mod == 0:
            for _ in range(13):
                change_per_day()
        else:
            for _ in range(mod-1):
                change_per_day()

        return cells