AdventOfCode2020/src/day18.py at main · KraProgrammer/AdventOfCode2020 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
from aocd.models import Puzzle
from lark import Lark, Transformer, v_args

calc_grammar = """
    ?start: mulsum
    ?mulsum: atom
        | mulsum "+" atom   -> add
        | mulsum "-" atom   -> sub
        | mulsum "*" atom  -> mul
        | mulsum "/" atom  -> div
    ?atom: NUMBER           -> number
         | NAME             -> var
         | "(" mulsum ")"
    %import common.CNAME -> NAME
    %import common.NUMBER
    %import common.WS_INLINE
    %ignore WS_INLINE
"""

calc_grammar2 = """
    ?start: product
    ?product: sum
        | sum "*" product  -> mul
        | sum "/" product  -> div
    ?sum: atom
        | sum "+" atom   -> add
        | sum "-" atom   -> sub
    ?atom: NUMBER           -> number
         | "(" product ")"
    %import common.CNAME -> NAME
    %import common.NUMBER
    %import common.WS_INLINE
    %ignore WS_INLINE
"""


@v_args(inline=True)  # Affects the signatures of the methods
class CalculateTree(Transformer):
    """
    Taken from https://github.com/lark-parser/lark/blob/master/examples/calc.py
    """
    number = float

    def __init__(self):
        super().__init__()
        self.vars = {}

    def add(self, mulsum, atom):
        return mulsum + atom

    def sub(self, mulsum, atom):
        return mulsum - atom

    def mul(self, mulsum, atom):
        return mulsum * atom

    def div(self, mulsum, atom):
        return mulsum / atom


def solve_puzzle_one(input_array):
    calc_parser = Lark(calc_grammar, parser='lalr', transformer=CalculateTree())
    calc = calc_parser.parse

    print(sum([calc(line) for line in input_array]))


def solve_puzzle_two(input_array):
    calc_parser = Lark(calc_grammar2, parser='lalr', transformer=CalculateTree())
    calc = calc_parser.parse

    print(sum([calc(line) for line in input_array]))


def parse_input(data):
    return data.splitlines()


test_input = """1 + 2 * 3 + 4 * 5 + 6
"""

test_input2 = """1 + (2 * 3) + (4 * (5 + 6))
"""

if __name__ == '__main__':
    puzzle = Puzzle(year=2020, day=18)
    input = puzzle.input_data
    if False:
        input = test_input2

    solve_puzzle_one(parse_input(input.strip()))
    solve_puzzle_two(parse_input(input.strip()))