aoc-2023/d09/main.py

89 lines
2.2 KiB
Python

from pathlib import Path
class DiffPyramid:
def __init__(self, starting: list[int]) -> None:
self.seq: list[list[int]] = [starting]
while not all([n == 0 for n in self.seq[-1]]):
self.seq.append(
[
self.seq[-1][i + 1] - self.seq[-1][i]
for i in range(0, len(self.seq[-1]) - 1)
]
)
def prepend(self, n: int) -> None:
for _ in range(n):
lasts = [s[0] for s in self.seq]
nexts: list[int] = []
for last in reversed(lasts):
if not nexts:
nexts.append(last)
else:
nexts.append(last - nexts[-1])
for i, next_num in enumerate(reversed(nexts)):
self.seq[i].insert(0, next_num)
def extend(self, n: int) -> None:
for _ in range(n):
lasts = [s[-1] for s in self.seq]
nexts: list[int] = []
for last in reversed(lasts):
if not nexts:
nexts.append(last)
else:
nexts.append(last + nexts[-1])
for i, next_num in enumerate(reversed(nexts)):
self.seq[i].append(next_num)
def __str__(self) -> str:
return "\n".join(
(" " * i) + " ".join(str(ss) for ss in s) for i, s in enumerate(self.seq)
)
def part1(input: Path) -> int:
total = 0
for line in input.read_text().split("\n"):
if not line:
break
pyramid = DiffPyramid([int(i) for i in line.split()])
pyramid.extend(1)
print(str(pyramid))
print()
total += pyramid.seq[0][-1]
return total
def part2(input: Path) -> int:
total = 0
for line in input.read_text().split("\n"):
if not line:
break
pyramid = DiffPyramid([int(i) for i in line.split()])
pyramid.prepend(1)
print(str(pyramid))
print()
total += pyramid.seq[0][0]
return total
if __name__ == "__main__":
# input = Path("sample.txt")
input = Path("input.txt")
result1 = part1(input)
print("part 1", result1)
result2 = part2(input)
print("part 2", result2)