advent-of-code/2021/python/day06.py

45 lines
1.1 KiB
Python

import shared
from collections import Counter, defaultdict
from pprint import pp
class Lanternfish:
def __init__(self, name):
self.load(name)
self.new = 0
def load(self, name):
with open(name, "r") as f:
self.nums = list(map(int, f.read().split(",")))
self.fish = defaultdict(int)
self.fish.update(dict(Counter(self.nums)))
def tick(self):
self.new = self.fish[0]
# could do for loop from 0-6 then handle 6 specially
self.fish[0] = self.fish[1]
self.fish[1] = self.fish[2]
self.fish[2] = self.fish[3]
self.fish[3] = self.fish[4]
self.fish[4] = self.fish[5]
self.fish[5] = self.fish[6]
self.fish[6] = self.fish[7] + self.new
self.fish[7] = self.fish[8]
self.fish[8] = self.new
def total(self):
return sum([v for k, v in self.fish.items() if k > -1])
def main():
l = Lanternfish(shared.get_fname(6))
TIMES = 256
for x in range(TIMES):
l.tick()
pp(l.total())
if __name__ == "__main__":
main()