Advent of Code 2019 Day 4, DP solution

advent-2019-4.py

"""
For part 1 for counting solutions ≤ x, our DP state can be dp[n][f][d][le] = the count of digit-strings with

- n digits
- where the first digit is f
- and a double digit has appeared iff d = 1
- that are less than equal to the last n digits of x iff le = 1.
"""

def part1_dp(x): # x should be list of digits
    # we'll waste dp[0][.][.][.] for clarity
    dp = [[[[0 for _ in range(2)] for _ in range(2)] for _ in range(10)] for _ in range(len(x) + 1)]

    for last_digit in range(10):
        dp[1][last_digit][0][last_digit <= x[-1]] = 1

    for digit_count in range(1, len(x)):
        for first_digit in range(10):
            for double_digit_exists in range(2):
                for le_x in range(2):
                    for next_digit in range(0, first_digit + 1):
                        next_double_digit_exists = double_digit_exists or next_digit == first_digit
                        x_digit = x[-digit_count-1]
                        next_le_x = next_digit < x_digit or (next_digit == x_digit and le_x)
                        dp[digit_count + 1][next_digit][next_double_digit_exists][next_le_x] += (
                            dp[digit_count][first_digit][double_digit_exists][le_x])

    return sum(dp[len(x)][first_digit][1][1] for first_digit in range(10))

def part1_solve(lo, hi):
    dlo = [int(d) for d in str(lo - 1)]
    dhi = [int(d) for d in str(hi)]
    return part1_dp(dhi) - part1_dp(dlo)

print(part1_solve(168630, 718098))

"""
part 2 for counting solutions ≤ x: our DP state can be dp[n][f][fr][d][le] = the count of digit-strings with

- n digits
- where the first digit is f
- it has repeated once if fr = 0, twice if fr = 1, or more if fr = 2
- a "safe" exactly-double-digit has appeared iff d = 1
- that are less than equal to the last n digits of x iff le = 1.
"""

def part2_dp(x):
    dp = [[[[[0 for _ in range(2)] for _ in range(2)] for _ in range(3)] for _ in range(10)] for _ in range(len(x) + 1)]

    for last_digit in range(10):
        dp[1][last_digit][0][0][last_digit <= x[-1]] = 1

    for digit_count in range(1, len(x)):
        for first_digit in range(10):
            for first_digit_repeats in range(3):
                for double_digit_exists in range(2):
                    for le_x in range(2):
                        for next_digit in range(0, first_digit + 1):
                            next_first_digit_repeats = min(first_digit_repeats + 1, 2) if next_digit == first_digit else 0
                            next_double_digit_exists = double_digit_exists or (next_digit != first_digit and first_digit_repeats == 1)
                            x_digit = x[-digit_count-1]
                            next_le_x = next_digit < x_digit or (next_digit == x_digit and le_x)
                            dp[digit_count + 1][next_digit][next_first_digit_repeats][next_double_digit_exists][next_le_x] += (
                                dp[digit_count][first_digit][first_digit_repeats][double_digit_exists][le_x])

    return sum(
            dp[len(x)][first_digit][0][1][1] +
            dp[len(x)][first_digit][1][1][1] +
            dp[len(x)][first_digit][2][1][1] +
            dp[len(x)][first_digit][1][0][1]
            for first_digit in range(10))

def part2_solve(lo, hi):
    dlo = [int(d) for d in str(lo - 1)]
    dhi = [int(d) for d in str(hi)]
    return part2_dp(dhi) - part2_dp(dlo)

print(part2_solve(168630, 718098))