day02: code beautification

This commit is contained in:
Stefan Harmuth 2021-12-02 06:29:31 +01:00
parent a0b4513d31
commit 2431b0fd8d

View File

@ -1,5 +1,21 @@
from aoc import AOCDay from aoc import AOCDay
from typing import Any from typing import Any, List
def follow_directions(path: List) -> (int, int, int):
pos = 0
depth = 0
aim = 0
for direction in path:
if direction[0] == 'forward':
pos += direction[1]
depth += aim * direction[1]
elif direction[0] == 'down':
aim += direction[1]
elif direction[0] == 'up':
aim -= direction[1]
return pos, depth, aim
class Day(AOCDay): class Day(AOCDay):
@ -8,30 +24,12 @@ class Day(AOCDay):
def part1(self) -> Any: def part1(self) -> Any:
path = self.getInputAsArraySplit(" ", [str, int]) path = self.getInputAsArraySplit(" ", [str, int])
pos = 0 pos, _, depth = follow_directions(path)
depth = 0
for direction in path:
if direction[0] == 'forward':
pos += direction[1]
elif direction[0] == 'down':
depth += direction[1]
elif direction[0] == 'up':
depth -= direction[1]
return pos * depth return pos * depth
def part2(self) -> Any: def part2(self) -> Any:
path = self.getInputAsArraySplit(" ", [str, int]) path = self.getInputAsArraySplit(" ", [str, int])
pos = 0 pos, depth, _ = follow_directions(path)
depth = 0
aim = 0
for direction in path:
if direction[0] == 'forward':
pos += direction[1]
depth += aim * direction[1]
elif direction[0] == 'down':
aim += direction[1]
elif direction[0] == 'up':
aim -= direction[1]
return pos * depth return pos * depth