AoC 2021 day 7: helping carbs conserve fuel

crabhelp.py

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

def pascal(distance):
    if distance<0: return pascal(-distance)
    if distance<2: return distance
    return (distance*(distance+1))//2

def move_all(xs, destination, func=abs):
    """
    returns amount of fuel to move all xs to that destination
    """
    return sum(func(destination-x) for x in xs)

def min_fuel(crab_xs, func=abs):
    # massively inefficient
    return min(move_all(crab_xs, pos, func) for pos in range(max(crab_xs)+1))

def part1(crab_xs):
    return min_fuel(crab_xs)

def part2(crab_xs):
    return min_fuel(crab_xs, func=pascal)

if __name__ == '__main__':
    with open('puzzle.dat') as f:
        crab_xs = list(map(int, f.readline().rstrip().split(",")))
    print("display result")
    print("part1", part1(crab_xs))
    print("part2", part2(crab_xs))

test.py

#!/usr/bin/env python3

import unittest

from crabhelp import *
    
with open('example.dat') as f:
    example_xs = list(map(int, f.readline().rstrip().split(",")))

class TestPart1(unittest.TestCase):
    def test_example(self):
        self.assertEqual( part1(example_xs), 37 ) # in a row?

class TestPart2(unittest.TestCase):
    def test_pascal(self):
        self.assertEqual( pascal(0), 0 )
        self.assertEqual( pascal(1), 1 )
        self.assertEqual( pascal(2), 3 )
        self.assertEqual( pascal(3), 6 )
        self.assertEqual( pascal(4), 10 )
        self.assertEqual( pascal(-3), 6 )
    def test_example(self):
        self.assertEqual( part2(example_xs), 168 )


if __name__ == '__main__':
    unittest.main()