Skip to content

957. Prison Cells After N Days 👎

  • Time:
  • Space: $O(1)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution {
 public:
  vector<int> prisonAfterNDays(vector<int>& cells, int n) {
    vector<int> firstDayCells;
    vector<int> nextDayCells(cells.size());

    for (int day = 0; n-- > 0; cells = nextDayCells, ++day) {
      for (int i = 1; i + 1 < cells.size(); ++i)
        nextDayCells[i] = cells[i - 1] == cells[i + 1];
      if (day == 0)
        firstDayCells = nextDayCells;
      else if (nextDayCells == firstDayCells)
        n %= day;
    }

    return cells;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution {
  public int[] prisonAfterNDays(int[] cells, int n) {
    int[] firstDayCells = new int[cells.length];
    int[] nextDayCells = new int[cells.length];

    for (int day = 0; n-- > 0; cells = nextDayCells.clone(), ++day) {
      for (int i = 1; i + 1 < cells.length; ++i)
        nextDayCells[i] = cells[i - 1] == cells[i + 1] ? 1 : 0;
      if (day == 0)
        firstDayCells = nextDayCells.clone();
      else if (Arrays.equals(nextDayCells, firstDayCells))
        n %= day;
    }

    return cells;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution:
  def prisonAfterNDays(self, cells: list[int], n: int) -> list[int]:
    nextDayCells = [0] * len(cells)
    day = 0

    while n > 0:
      n -= 1
      for i in range(1, len(cells) - 1):
        nextDayCells[i] = 1 if cells[i - 1] == cells[i + 1] else 0
      if day == 0:
        firstDayCells = nextDayCells.copy()
      elif nextDayCells == firstDayCells:
        n %= day
      cells = nextDayCells.copy()
      day += 1

    return cells