# Leetcode 957. Prison Cells After N Days

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

• You need to de-indent this code. – Reinderien Jul 6 at 18:08
• 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). – rdllopes Jul 6 at 18:38

## 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()
while True:
next_node = tuple(step_map[previous])
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).