Advent of Code 2020 Day 9 in Python

Question

This is a solution for part 1 & 2 of Advent of Code 2020 Day 9.

Given the example input sample.txt:

35
20
15
25
47
40
62
55
65
95
102
117
150
182
127
219
299
277
309
576

Part 1: find the first number in the list (after the preamble) which is not the sum of two of the n numbers before it. n is the preamble.

For example, after the 5-number preamble, almost every number is the sum of two of the previous 5 numbers; the only number that does not follow this rule is 127.

Part 2: find a contiguous set of at least two numbers in your list which sum to the invalid number from part 1.

For example, in this list, adding up all of the numbers from 15 through 40 produces the invalid number from part 1, 127. Add together the smallest and largest number in this contiguous range; in this example, these are 15 and 47, producing 62.

Code:

def not_in_preamble(nums, preamble):
    diff = set()
    for i in range(preamble, len(nums)):
        current = nums[i]
        found = False
        for j in range(i - preamble, i):
            if nums[j] in diff:
                found = True
                break
            diff.add(current - nums[j])
        if not found:
            return current
        diff.clear()


def sub_array_with_sum(nums, n):
    i, j = 0, 1
    partial_sum = nums[i] + nums[j]
    while j < len(nums):
        if partial_sum == n:
            return nums[i: j + 1]
        if partial_sum < n or i == j - 1:
            j += 1
            partial_sum += nums[j]
        else:
            partial_sum -= nums[i]
            i += 1


# Part 1
preamble = 5
with open('sample.txt') as f:
    nums = [int(l) for l in f.read().splitlines()]
pt1_result = not_in_preamble(nums, preamble)
print(f'PT 1 result: {pt1_result}')

# Part 2
n = pt1_result
sub_array = sub_array_with_sum(nums, n)
print(f'PT 2 result {min(sub_array) + max(sub_array)}')

The code outputs:

PT 1 result: 127
PT 2 result 62

Which is the correct result for the sample input.

Changing the sample.txt to the puzzle input and setting preamble=25, will give the correct result for the challenge.

Can the code be made more compact and efficient? Any other feedback or alternative solutions are welcome.

Answer 1

Part 1

For not_in_preamble, I'd use a for-else instead of your found variable, and not re-use the same set (I don't think it's worth it):

def not_in_preamble(nums, preamble):
    for i in range(preamble, len(nums)):
        current = nums[i]
        diff = set()
        for j in range(i - preamble, i):
            if nums[j] in diff:
                break
            diff.add(current - nums[j])
        else:
            return current

Since your "Any other feedback or alternative solutions are welcome" also allows less efficient solutions, here's how I actually solved it myself :-)

from itertools import combinations

def not_in_preamble(nums, preamble):
    for i in range(preamble, len(nums)):
        if nums[i] not in map(sum, combinations(nums[i-preamble : i], 2)):
            return nums[i]

Preamble size 25 and nums size 1000 are easily small enough for that to run quickly.

Part 2

Your sub_array_with_sum is great, although it does rely on the input being non-negative numbers, which isn't guaranteed in the problem statement (but is the case for the fixed input file).

I'd just go with PEP 8 and change nums[i: j + 1] to nums[i : j + 1] or nums[i : j+1]. And perhaps I'd go with while True:, relying on the existence of an answer.

Your algorithm almost only takes O(1) space. The only reason it takes more is copying out the sub-array (or rather sub-list, as this is Python and the problem never says "array", so there's really no need to call it array). Could be done in O(1) space to make the algorithm not just time-optimal (O(n)) but also space-optimal. Though that would require more code, not worth it (unless you have an interviewer who insists).

Input

I'd change [int(l) for l in f.read().splitlines()] to list(map(int, f)), or at least [int(line) for line in f].

Output

I'd change
print(f'PT 1 result: {pt1_result}') to
print('PT 1 result:', pt1_result). Bit shorter, and it works in older versions of Python. And at least in IDLE, the f-string is styled only as string, the whole thing green, including the variable. I don't like that.