88 lines
2.2 KiB
Python
88 lines
2.2 KiB
Python
from lib import get_data, str_to_ints
|
|
from math import lcm
|
|
|
|
|
|
def freeze(coords, velos):
|
|
coords = tuple(map(tuple, coords))
|
|
velos = tuple(map(tuple, velos))
|
|
return coords + velos
|
|
|
|
|
|
def part_1(data):
|
|
coords = []
|
|
velos = []
|
|
for line in data.splitlines():
|
|
coords.append(str_to_ints(line))
|
|
velos.append([0, 0, 0])
|
|
|
|
for _ in range(1_000):
|
|
# apply gravity
|
|
for i in range(len(coords)):
|
|
for j in range(i + 1, len(coords)):
|
|
for d in range(3):
|
|
if coords[i][d] > coords[j][d]:
|
|
velos[i][d] -= 1
|
|
velos[j][d] += 1
|
|
elif coords[i][d] < coords[j][d]:
|
|
velos[i][d] += 1
|
|
velos[j][d] -= 1
|
|
|
|
# update coords
|
|
for i in range(len(coords)):
|
|
for d in range(3):
|
|
coords[i][d] += velos[i][d]
|
|
|
|
r = 0
|
|
for i in range(len(coords)):
|
|
p = sum(map(abs, coords[i]))
|
|
k = sum(map(abs, velos[i]))
|
|
r += p * k
|
|
print(r)
|
|
|
|
|
|
def part_2(data):
|
|
steps = []
|
|
for i in range(3):
|
|
coords = []
|
|
velos = []
|
|
|
|
for line in data.splitlines():
|
|
coords.append(str_to_ints(line)[i])
|
|
velos.append(0)
|
|
|
|
seen = {}
|
|
for step in range(1_000_000):
|
|
state = tuple(coords) + tuple(velos)
|
|
if state in seen:
|
|
steps.append(step)
|
|
break
|
|
else:
|
|
seen[state] = step
|
|
|
|
for i in range(len(coords)):
|
|
for j in range(i + 1, len(coords)):
|
|
if coords[i] > coords[j]:
|
|
velos[i] -= 1
|
|
velos[j] += 1
|
|
elif coords[i] < coords[j]:
|
|
velos[i] += 1
|
|
velos[j] -= 1
|
|
|
|
# update coords
|
|
for i in range(len(coords)):
|
|
coords[i] += velos[i]
|
|
|
|
# Intuition: Find when position repeats on each axis and then find lowest
|
|
# common multiple for all three values.
|
|
print(lcm(*steps))
|
|
|
|
|
|
def main():
|
|
data = get_data(__file__)
|
|
part_1(data)
|
|
part_2(data)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|