LeetCode 957: Prison Cells After N Days

Problem Restatement

We are given an array cells of length 8.

Each cell is either:

0  # vacant
1  # occupied

Every day, the cells change using this rule:

If a cell has two adjacent neighbors that are both occupied or both vacant, the cell becomes occupied. Otherwise, it becomes vacant.

The first and last cells do not have two neighbors, so they always become vacant after one day.

Given an integer n, return the prison state after n days.

The official constraints are:

Source: LeetCode problem statement.

Input and Output

Example function shape:

def prisonAfterNDays(cells: list[int], n: int) -> list[int]:
    ...

Examples

Example 1:

cells = [0, 1, 0, 1, 1, 0, 0, 1]
n = 7

Output:

[0, 0, 1, 1, 0, 0, 0, 0]

The state changes day by day until day 7.

Example 2:

cells = [1, 0, 0, 1, 0, 0, 1, 0]
n = 1000000000

Output:

[0, 0, 1, 1, 1, 1, 1, 0]

Here n is very large, so simulating every day directly is too slow.

First Thought: Simulate Every Day

The transition for one day is simple.

For each middle cell i, where 1 <= i <= 6, compute:

next_cells[i] = 1 if cells[i - 1] == cells[i + 1] else 0

The first and last cells are always 0:

next_cells[0] = 0
next_cells[7] = 0

A direct simulation would repeat this process n times.

That works for small n, but n can be as large as 10^9, so direct simulation may be too slow.

Key Insight

There are only 8 cells, and each cell has two possible states.

So there are at most:

2 ** 8 = 256

possible states.

After the first day, the first and last cells are always 0, so only the middle 6 cells can vary. That gives at most:

2 ** 6 = 64

states.

If we keep simulating, eventually we must see a state that we have seen before. Once a state repeats, the future states repeat in the same cycle.

So instead of simulating all n days, we detect the cycle and skip full repetitions.

Algorithm

Use a dictionary:

seen[state] = remaining_days

where state is the current cell state as a tuple.

Then:

1. While n > 0:
   • Convert cells to a tuple.
   • If this state was seen before:
     • Compute the cycle length.
     • Reduce n using modulo.
   • Store the current state with the current n.
   • If n > 0, simulate one day.
   • Decrease n by 1.
2. Return cells.

The cycle length is:

cycle_length = seen[state] - n

because seen[state] records how many days were left the previous time we saw the same state.

Then we can skip full cycles:

n %= cycle_length

Correctness

The transition from one day to the next depends only on the current state. It does not depend on how we reached that state.

Therefore, if the same state appears twice, the sequence of future states from that point onward must also repeat.

The algorithm records every visited state. When it sees a repeated state, it knows the cycle length exactly: the number of simulated days between the first occurrence and the second occurrence.

Reducing n modulo the cycle length is safe because completing one full cycle returns to the same state. Removing any number of full cycles does not change the final state.

After cycle reduction, the algorithm simulates only the remaining days. Each simulation step applies the exact transition rule from the problem. Therefore, the returned state is exactly the prison state after n days.

Complexity

Because there are only 8 cells, the number of possible states is constant.

More explicitly, the algorithm simulates at most 256 states, and after day one at most 64 meaningful states.

Implementation

class Solution:
    def prisonAfterNDays(self, cells: list[int], n: int) -> list[int]:
        seen = {}

        while n > 0:
            state = tuple(cells)

            if state in seen:
                cycle_length = seen[state] - n
                n %= cycle_length

            seen[state] = n

            if n > 0:
                n -= 1
                cells = self.next_day(cells)

        return cells

    def next_day(self, cells: list[int]) -> list[int]:
        next_cells = [0] * 8

        for i in range(1, 7):
            if cells[i - 1] == cells[i + 1]:
                next_cells[i] = 1

        return next_cells

Code Explanation

We store seen states in a dictionary:

seen = {}

Since lists are mutable and cannot be dictionary keys, we convert the current state to a tuple:

state = tuple(cells)

If the state appeared before:

if state in seen:

we compute the cycle length:

cycle_length = seen[state] - n

Then we skip all full cycles:

n %= cycle_length

We store the current state with the current number of remaining days:

seen[state] = n

Then, if there are still days left, we simulate one day:

if n > 0:
    n -= 1
    cells = self.next_day(cells)

The helper function builds the next state:

next_cells = [0] * 8

This automatically sets the first and last cells to 0.

For middle cells:

if cells[i - 1] == cells[i + 1]:
    next_cells[i] = 1

A cell becomes occupied when its two neighbors are equal.

Testing

def run_tests():
    s = Solution()

    assert s.prisonAfterNDays(
        [0, 1, 0, 1, 1, 0, 0, 1],
        7,
    ) == [0, 0, 1, 1, 0, 0, 0, 0]

    assert s.prisonAfterNDays(
        [1, 0, 0, 1, 0, 0, 1, 0],
        1000000000,
    ) == [0, 0, 1, 1, 1, 1, 1, 0]

    assert s.prisonAfterNDays(
        [0, 0, 0, 0, 0, 0, 0, 0],
        1,
    ) == [0, 1, 1, 1, 1, 1, 1, 0]

    assert s.prisonAfterNDays(
        [1, 1, 1, 1, 1, 1, 1, 1],
        1,
    ) == [0, 1, 1, 1, 1, 1, 1, 0]

    assert s.prisonAfterNDays(
        [1, 0, 1, 0, 1, 0, 1, 0],
        1,
    ) == [0, 1, 1, 1, 1, 1, 1, 0]

    print("all tests passed")

run_tests()

Test meaning: