part2_func.py 642 B

12345678910111213141516171819
  1. import argparse
  2. from typing import List, Set, NamedTuple
  3. import functools
  4. import operator
  5. parser = argparse.ArgumentParser()
  6. parser.add_argument("ifile", type=argparse.FileType('r'))
  7. args = parser.parse_args()
  8. fish = [int(x) for x in args.ifile.readline().split(',')]
  9. new_fish_per_day = [0] * (255)
  10. new_fish_per_day = list(map(sum, zip(*[map(lambda x: int((x % 7) == 0), range(start, -len(new_fish_per_day), -1)) for start in fish])))
  11. for index, num_new_fish in enumerate(new_fish_per_day):
  12. for i in range(index + 1 + 8, len(new_fish_per_day), 7):
  13. new_fish_per_day[i] += num_new_fish
  14. print(sum(new_fish_per_day) + len(fish))