Program to find state of prison cells after k days in python


Suppose we have a binary list (1s and 0s in the list) and another value k. Each value in nums represents the state of a prison cell where 1 indicates occupied cell and 0 indicates empty cell. On each day when a cell has two adjacent cells that are either both occupied or both empty, then it becomes occupied. Otherwise, it becomes empty. So we have to find state of the prison cells after k number of days.

So, if the input is like nums = [1, 0, 1, 0, 0, 0, 0, 0] k = 1, then the output will be [0, 1, 1, 0, 1, 1, 1, 0], as we can notice the first and last index can never be occupied because they can never have 2 neighbors.

To solve this, we will follow these steps:

  • Define a function next_day_state() . This will take cells
  • new_cells := a copy of cells
  • new_cells[0] := 0, new_cells[7] := 0
  • for j in range 1 to 6, do
    • if cells[j - 1] is same as cells[j + 1], then
      • new_cells[j] := 1
    • otherwise,
      • new_cells[j] := 0
  • return new_cells
  • From the main method do the following:
  • seen := a new map
  • flag := False, i := 0
  • while i < N, do
    • ns := next_day_state(cells)
    • if ns is not seen, then
      • mark ns as seen
    • otherwise,
      • flag := True
      • come out from the loop
    • cells := ns
    • i := i + 1
  • if flag is true, then
    • N := N mod (number of items seen)
    • i := 0
    • while i < N is non-zero, do
      • ns := next_day_state(cells)
      • i := i + 1
      • cells := ns
  • return cells

Let us see the following implementation to get better understanding:

Example

Live Demo

import copy
class Solution:
   def next_day_state(self, cells):
      new_cells = copy.copy(cells)
      new_cells[0] = 0
      new_cells[7] = 0
      for j in range(1, 7):
         if cells[j - 1] == cells[j + 1]:
            new_cells[j] = 1
         else:
            new_cells[j] = 0
      return new_cells

   def solve(self, cells, N):
      seen = dict()
      flag, i = False, 0

      while i < N:
         ns = self.next_day_state(cells)
         if tuple(ns) not in seen:
            seen[tuple(ns)] = True
         else:
            flag = True
            break
         cells = ns
         i += 1

      if flag:
         N = N % len(seen)
         i = 0
         while i < N:
            ns = self.next_day_state(cells)
            i += 1
            cells = ns
      return cells

ob = Solution()
nums = [1, 0, 1, 0, 0, 0, 0, 0]
k = 1
print(ob.solve(nums, k))

Input

[4, 7, 2, 5], 6

Output

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

Updated on: 26-Nov-2020

302 Views

Kickstart Your Career

Get certified by completing the course

Get Started
Advertisements