zerosleeps

Since 2010

Advent of Code 2020 day 1

My Python solution for Advent of Code 2020 day 1. I’ve cleaned this up a wee bit since submitting my answers, but the essence of it hasn’t changed. I did use itertools.permutations at first, which happened to work because the loop bails out as soon as it finds a solution, but I changed it to itertools.combinations for readability.

from pathlib import Path
from itertools import combinations
from math import prod
import unittest

def get_input():
    input_file = Path(__file__).parent / 'day_01_input.txt'

    with open(input_file) as file:
        return [int(line.strip()) for line in file.readlines()]

def run(input, length):
    for p in combinations(input, length):
        if sum(p) == 2020:
            return(prod(p))

class TestExamples(unittest.TestCase):
    def setUp(self):
        self.input = [1721,979,366,299,675,1456]

    def test_part_one(self):
        self.assertEqual(run(self.input, 2), 514579)

    def test_part_two(self):
        self.assertEqual(run(self.input, 3), 241861950)

if __name__ == '__main__':
    print(f'Part one: {run(get_input(), 2)}')
    print(f'Part two: {run(get_input(), 3)}')