adventofcode2022/day20/sol.py

data = []
for line in open("input"):
    data.append(int(line.strip()))

#data = [1, 2, -3, 3, -2, 0, 4, ]


def mix(data, rounds=1):
    #assert len(set(data)) == len(data)
    idx = list(range(len(data)))
    N = len(data) - 1

    for _ in range(rounds):

        for k in range(len(data)):
            i = idx.index(k)
            x = data[k]
            if x != 0:
                j = (i+x) % N
                if x < 0 and j == 0:
                    j = N
                idx.insert(j, idx.pop(i))
            else:
                j = i
            #print(x, i, j, [data[k] for k in idx])

    return [data[k] for k in idx]

def answer(data):
    i = data.index(0)
    n = []
    for j in 1000, 2000, 3000:
        n.append( data[(i+j)%len(data)] )
    print(n)
    return sum(n)

#print(data)

copy = mix(data)
print(answer(copy))


p2 = [x*811589153 for x in data]
print(answer(mix(p2, 10)))