-
Notifications
You must be signed in to change notification settings - Fork 1
/
aoc202014.py
116 lines (88 loc) · 3.07 KB
/
aoc202014.py
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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
"""AoC 14, 2020: Docking Data."""
# Standard library imports
import pathlib
import sys
# Third party imports
import parse
MASK_PATTERN = parse.compile("mask = {mask}")
MEMORY_PATTERN = parse.compile("mem[{address:d}] = {value:d}")
def parse_data(puzzle_input):
"""Parse input."""
program = []
for line in puzzle_input.split("\n"):
if match := MASK_PATTERN.parse(line):
active_mask = match["mask"]
if match := MEMORY_PATTERN.parse(line):
program.append((active_mask, match["address"], match["value"]))
return program
def part1(data):
"""Solve part 1."""
memory = {
address: apply_mask_to_value(mask, value) for mask, address, value in data
}
return sum(memory.values())
def part2(data):
"""Solve part 2."""
memory = {
address: value
for mask, address_seed, value in data
for address in construct_masks(update_mask_with_address(mask, address_seed))
}
return sum(memory.values())
def apply_mask_to_value(mask, value):
"""Apply mask to value.
The current bitmask is applied to values immediately before they are written
to memory: a 0 or 1 overwrites the corresponding bit in the value, while an
X leaves the bit in the value unchanged.
## Examples:
>>> apply_mask_to_value("XXXXXXXXXXX10XXXX1", 42)
75
"""
num_bits = len(mask)
value_str = ("0" * num_bits + bin(value)[2:])[-num_bits:]
return int(
"".join(
onoff if mask_char == "X" else mask_char
for onoff, mask_char in zip(value_str, mask)
),
base=2,
)
def construct_masks(mask):
"""Find all possible masks by resolving floating bits.
If the bitmask bit is X, the corresponding memory address bit is floating.
A floating bit is not connected to anything and instead fluctuates
unpredictably. In practice, this means the floating bits will take on all
possible values, potentially causing many memory addresses to be written all at
once!
## Example:
>>> list(construct_masks("00X1X"))
['00010', '00011', '00110', '00111']
"""
idx = mask.find("X")
if idx == -1:
yield mask
else:
yield from construct_masks(f"{mask[:idx]}0{mask[idx + 1:]}")
yield from construct_masks(f"{mask[:idx]}1{mask[idx + 1:]}")
def update_mask_with_address(mask, address):
"""Apply mask to address.
## Example:
>>> update_mask_with_address("0X0X100X", 42)
'0X1X101X'
"""
num_bits = len(mask)
address_str = ("0" * num_bits + bin(address)[2:])[-num_bits:]
return "".join(
mask_char if mask_char in {"X", "1"} else onoff
for onoff, mask_char in zip(address_str, mask)
)
def solve(puzzle_input):
"""Solve the puzzle for the given input."""
data = parse_data(puzzle_input)
yield part1(data)
yield part2(data)
if __name__ == "__main__":
for path in sys.argv[1:]:
print(f"\n{path}:")
solutions = solve(puzzle_input=pathlib.Path(path).read_text().strip())
print("\n".join(str(solution) for solution in solutions))