Leetcode 957. Prison Cells After N Days

Question

I am working on 957. Prison Cells After N Days but getting timeout error for my code. as I assume my code should be 14+O(n-14) about O(n) am I right since after 14 run pattern will be repeated? and how can I improve my code?

 from collections import defaultdict
 step_map = defaultdict(list)
 for k in range(N):
 if tuple(cells) in step_map:
 cells = step_map[tuple(cells)]
 else:
 tmp = list()
 for i in range(1, 7):
 tmp += [int(cells[i - 1] == cells[i + 1])]
 tmp = [0] + tmp + [0]
 step_map[tuple(cells)] = tmp
 cells = tmp
 return cells

• Reinderien: You need to de-indent this code.
• rdllopes: the complexity time is 8*N*k (there are 8 cells not 7). in Big-O notation, it does not make sense to say 16 + O(n - 16)... it should be O(n).

Answer 1

The important part

They want something below O(n). Using map is a good idea but actually you should find the cycles and return the right position on the cycle instead of computing line by line.

Spoiler (possible solution)

Change:

 if tuple(cells) in step_map:
 return step_map[tuple(cells)]

to:

 if tuple(cells) in step_map:
 cycle = list()
 head = tuple(cells)
 cycle.append(head)
 previous = head
 while True:
 next_node = tuple(step_map[previous])
 if next_node == head:
 return list(cycle[(N - k) % len(cycle)])
 cycle.append(next_node)
 previous = next_node

---

Old edit - Some small improvements

There are some O(m) operations multiple times...

For example:

 tmp = [0] + tmp + [0]

Python operation for that is O(m). Therefore, your solution is O(nm).

step_map[tuple(cells)] = tmp

this is also O(m).