forked from lxylxy123456/AdventOfCode2024
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy paths.py
More file actions
81 lines (72 loc) · 1.78 KB
/
s.py
File metadata and controls
81 lines (72 loc) · 1.78 KB
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
# Youtube: https://youtu.be/jdkzd5PXo4A
import argparse, math, sys, re, functools, operator, itertools, heapq
from collections import defaultdict, Counter, deque
#sys.setrecursionlimit(100000000)
#A = list(map(int, input().split()))
#T = int(input())
def read_lines(f):
while True:
line = f.readline()
if not line:
break
assert line[-1] == '\n'
yield line[:-1]
def main():
parser = argparse.ArgumentParser()
parser.add_argument('-1', '--one', action='store_true', help='Only part 1')
parser.add_argument('-2', '--two', action='store_true', help='Only part 2')
parser.add_argument('input_file', nargs='?')
args = parser.parse_args()
if args.input_file is not None:
f = open(args.input_file)
else:
f = sys.stdin
lines = list(read_lines(f))
if not args.two:
print(part_1(lines))
if not args.one:
print(part_2(lines))
@functools.lru_cache
def check_equation1(lhs, rhs):
if len(rhs) == 1:
return lhs == rhs[0]
# +
if check_equation1(lhs, (rhs[0] + rhs[1], *rhs[2:])):
return True
# *
if check_equation1(lhs, (rhs[0] * rhs[1], *rhs[2:])):
return True
return False
def part_1(lines):
s = 0
for i in lines:
lhs, rhs = i.split(':')
lhs = int(lhs)
rhs = tuple(map(int, rhs.split()))
if check_equation1(lhs, rhs):
s += lhs
return s
def check_equation2(lhs, rhs):
if len(rhs) == 1:
return lhs == rhs[0]
# +
if check_equation2(lhs, (rhs[0] + rhs[1], *rhs[2:])):
return True
# *
if check_equation2(lhs, (rhs[0] * rhs[1], *rhs[2:])):
return True
# ||
if check_equation2(lhs, (int(str(rhs[0]) + str(rhs[1])), *rhs[2:])):
return True
return False
def part_2(lines):
s = 0
for i in lines:
lhs, rhs = i.split(':')
lhs = int(lhs)
rhs = tuple(map(int, rhs.split()))
if check_equation2(lhs, rhs):
s += lhs
return s
if __name__ == '__main__':
main()