AoC 2021 day 6: lanternfish

aoc2021d6.py

#!/usr/bin/env python3
#
# https://adventofcode.com/2021/day/6

from collections import defaultdict

class School:
    def __init__(self, allfish, spawntime=7, firstspawntime=8):
        if type(allfish) is str: # assume filename
            with open(allfish) as fi:
                allfish = list(int(item) 
                        for item in fi.readline().rstrip().split(","))
        else: #already a list
            allfish = list(allfish)
        self.spawntime = spawntime
        self.firstspawntime = firstspawntime
        self.times = defaultdict(int)
        self.times = {time: allfish.count(time) 
            for time in range(firstspawntime)}
    def __len__(self):
        return sum(self.times.values())
    def step(self, ct=1):
        oldtimes = self.times
        for i in range(ct):
            newtimes = {time-1: oldtimes.get(time,0)
                    for time in range(1,firstspawntime+1)} 
            newtimes[firstspawntime] = oldtimes.get(0,0)
            newtimes[spawntime-1] += oldtimes.get(0,0)
            oldtimes=newtimes
        self.times = newtimes
    def count(self, i):
        return self.times[i]
    def __len__(self):
        return sum(self.times.values())

def assertEqual(actual,expected):
    assert actual == expected, "%r should equal %r" % (actual,expected)

testSchool = School([3])
assertEqual( len(testSchool), 1 )
testSchool.step()
assertEqual( testSchool.count(3), 0)
assertEqual( testSchool.count(2), 1)
testSchool.step()
assertEqual( testSchool.count(2), 0)
assertEqual( testSchool.count(1), 1)
testSchool.step()
assertEqual( testSchool.count(1), 0)
assertEqual( testSchool.count(0), 1)
testSchool.step()
assertEqual( testSchool.count(0), 0)
assertEqual( testSchool.count(6), 1)
assertEqual( testSchool.count(8), 1)
testSchool.step()
assertEqual( testSchool.count(0), 0)
assertEqual( testSchool.count(5), 1)
assertEqual( testSchool.count(7), 1)
assertEqual( len(testSchool), 2 )
# test multi
testSchool = School([3])
assertEqual( len(testSchool), 1 )
testSchool.step(5)
assertEqual( testSchool.count(0), 0)
assertEqual( testSchool.count(5), 1)
assertEqual( testSchool.count(7), 1)
assertEqual( len(testSchool), 2 )

d6test = School("d6test.txt")
d6test.step(18)
assertEqual( len(d6test), 26 )
d6test.step(62)
assertEqual( len(d6test), 5934 )

d6data = School("d6data.txt")
d6data.step(80)
print("part1:", len(d6data))

d6test.step( 256-80 )
assertEqual( len(d6test), 26984457539 )
d6data.step( 256-80 )
print("part2:", len(d6data))