diff --git a/aoc2020/day01/solution.py b/aoc2020/day01/solution.py index 739e493..2e08da4 100644 --- a/aoc2020/day01/solution.py +++ b/aoc2020/day01/solution.py @@ -3,17 +3,46 @@ import typing def main() -> int: - i, j = _find_2020_pair([int(s) for s in sys.stdin.read().split()]) - print("{} * {} == {}".format(i, j, i * j)) + inputs = [int(s) for s in sys.stdin.read().split()] + for i, j in _find_2020_pairs(inputs): + print(f"pair: {i} * {j} == {i * j}") + + for i, j, k in _find_2020_triplets(inputs): + print(f"triplet: {i} * {j} * {k} == {i * j * k}") return 0 -def _find_2020_pair(inputs: typing.List[int]) -> typing.Tuple[int, int]: - for i in inputs: - for j in inputs: - if i + j == 2020: - return i, j - return 0, 0 +def _find_2020_pairs( + inputs: typing.List[int], +) -> typing.Generator[None, typing.Tuple[int, int], None]: + found = [] + + for i, in0 in enumerate(inputs): + for j, in1 in enumerate(inputs): + if i == j: + continue + if in0 + in1 == 2020: + to_yield = tuple(sorted([in0, in1])) + if to_yield not in found: + yield to_yield + found.append(to_yield) + + +def _find_2020_triplets( + inputs: typing.List[int], +) -> typing.Generator[None, typing.Tuple[int, int], None]: + found = [] + + for i, in0 in enumerate(inputs): + for j, in1 in enumerate(inputs): + for k, in2 in enumerate(inputs): + if i == j or i == k: + continue + if in0 + in1 + in2 == 2020: + to_yield = tuple(sorted([in0, in1, in2])) + if to_yield not in found: + yield to_yield + found.append(to_yield) if __name__ == "__main__":