kot.pink
1import sys
2from operator import mul, add
3from functools import reduce
4
5total_wordwise = 0
6total_digitwise = 0
7file = '6.input' if len(sys.argv) <= 1 else sys.argv[1]
8
9columns = []
10ops = []
11op_line = None
12
13for line in open(file).readlines():
14 line = line.strip('\n')
15 for word in line.split(' '):
16 if not word: continue
17 if word in '+*':
18 ops.append(word)
19 op_line = line
20 else:
21 columns.append(word)
22
23for op_i, op in enumerate(ops):
24 column = [int(c) for i, c in enumerate(columns) if i % len(ops) == op_i]
25 if op == '+':
26 total_wordwise += sum(column)
27 elif op == '*':
28 total_wordwise += reduce(mul, column, 1)
29
30print(f'p1: {total_wordwise}')
31
32lines = []
33op_pos = [i for i, ch in enumerate(op_line) if ch in '+*']
34for line in open(file).readlines():
35 if line.strip() == op_line.strip(): break
36 words = [line[i:j][:-1] for i, j in zip([0] + op_pos[1:], op_pos[1:] + [None])]
37 lines.append(words)
38
39for op_i, op in enumerate(ops):
40 cols = [line[i] for i in range(len(lines[0])) if i % len(ops) == op_i for line in lines]
41 digits = [col[i] for i in range(len(cols[0])) for col in cols]
42 numbers = [int(''.join(digits[i:i+len(cols)])) for i in range(0, len(digits), len(cols))]
43 if op == '+':
44 total_digitwise += sum(numbers)
45 elif op == '*':
46 total_digitwise += reduce(mul, numbers, 1)
47
48print(f'p2: {total_digitwise}')