advent_fish.py

#for k in [5,1,2,1,5,3,1,1,1,1,1,2,5,4,1,1,1,1,2,1,2,1,1,1,1,1,2,1,5,1,1,1,3,1,1,1,3,1,1,3,1,1,4,3,1,1,4,1,1,1,1,2,1,1,1,5,1,1,5,1,1,1,4,4,2,5,1,1,5,1,1,2,2,1,2,1,1,5,3,1,2,1,1,3,1,4,3,3,1,1,3,1,5,1,1,3,1,1,4,4,1,1,1,5,1,1,1,4,4,1,3,1,4,1,1,4,5,1,1,1,4,3,1,4,1,1,4,4,3,5,1,2,2,1,2,2,1,1,1,2,1,1,1,4,1,1,3,1,1,2,1,4,1,1,1,1,1,1,1,1,2,2,1,1,5,5,1,1,1,5,1,1,1,1,5,1,3,2,1,1,5,2,3,1,2,2,2,5,1,1,3,1,1,1,5,1,4,1,1,1,3,2,1,3,3,1,3,1,1,1,1,1,1,1,2,3,1,5,1,4,1,3,5,1,1,1,2,2,1,1,1,1,5,4,1,1,3,1,2,4,2,1,1,3,5,1,1,1,3,1,1,1,5,1,1,1,1,1,3,1,1,1,4,1,1,1,1,2,2,1,1,1,1,5,3,1,2,3,4,1,1,5,1,2,4,2,1,1,1,2,1,1,1,1,1,1,1,4,1,5]:
#    fish[k] += 1

"""
Debug stuff



def invert_ind_at_step(step, ind):
    shift = (ind - step) % SIZE
    if shift >= 0:
        w = shift
    else:
        w = SIZE + shift
    return w

def unshifted_fish(n):
    return list(fish[ind_at_step(n, r)] for r in range(0, SIZE))
"""

SIZE = 9
FORK_SRC = 0
FORK_DST = 7
SOURCE_FISH = [3,4,3,1,2]

fish = [0] * SIZE
for k in SOURCE_FISH:
    fish[k] += 1

def ind_at_step(step, ind):
    return (ind + step) % SIZE

for n in range(0, 256):
    fish[ind_at_step(n, FORK_DST)] += fish[ind_at_step(n, FORK_SRC)]

print (sum(fish))