-
Notifications
You must be signed in to change notification settings - Fork 0
/
part_a.py
executable file
·143 lines (125 loc) · 4.85 KB
/
part_a.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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
#!/usr/bin/env python3
import re
from dataclasses import dataclass, field
from typing import Dict, List, Tuple, Iterable, Set
from aox.challenge import Debugger
from utils import BaseChallenge, helper
class Challenge(BaseChallenge):
def solve(self, _input, debugger: Debugger):
"""
>>> Challenge().default_solve()
518
"""
machine, start = Machine.from_substitutions_text_and_start(_input)
return machine.get_all_possible_distinct_next_step_count(start)
@dataclass
class Machine:
substitutions: Dict[str, List[str]] = field(default_factory=dict)
re_substitution = re.compile(r"^(\w+) => (\w+)$")
@classmethod
def from_substitutions_text_and_start(
cls, substitutions_and_start_text: str):
"""
>>> Machine.from_substitutions_text_and_start(
... "H => HO\\n"
... "H => OH\\n"
... "O => HH\\n"
... "\\n"
... "HOH\\n"
... )
(Machine(substitutions={'H': ['HO', 'OH'], 'O': ['HH']}), 'HOH')
"""
substitutions_text, start = \
cls.split_substitutions_and_start_text(substitutions_and_start_text)
return cls.from_substitutions_text(substitutions_text), start
@classmethod
def split_substitutions_and_start_text(
cls, substitutions_and_start_text: str) -> Tuple[str, str]:
"""
>>> Machine.split_substitutions_and_start_text(
... "H => HO\\nH => OH\\nO => HH\\n\\nHOH\\n")
('H => HO\\nH => OH\\nO => HH', 'HOH')
"""
substitutions_text, start = substitutions_and_start_text\
.strip().split('\n\n')
return substitutions_text, start
@classmethod
def from_substitutions_text(cls, substitutions_text: str):
"""
>>> Machine.from_substitutions_text(
... "H => HO\\n"
... "H => OH\\n"
... "O => HH\\n"
... )
Machine(substitutions={'H': ['HO', 'OH'], 'O': ['HH']})
"""
substitution_items = map(
cls.parse_substitution, substitutions_text.splitlines())
return cls.from_substitution_items(substitution_items)
@classmethod
def from_substitution_items(
cls, substitution_items: Iterable[Tuple[str, str]]):
"""
>>> Machine.from_substitution_items([
... ("H", "HO"),
... ("H", "OH"),
... ("O", "HH"),
... ])
Machine(substitutions={'H': ['HO', 'OH'], 'O': ['HH']})
"""
return cls(helper.group_by(
substitution_items, key=lambda pair: pair[0],
value=lambda pair: pair[1]))
@classmethod
def parse_substitution(cls, substitution_text: str) -> Tuple[str, str]:
match = cls.re_substitution.match(substitution_text)
if not match:
raise Exception(
f"Could not parse substitution {repr(substitution_text)}")
start, result = match.groups()
return start, result
def get_all_possible_distinct_next_step_count(self, start: str) -> int:
"""
>>> Machine({'H': ['HO', 'OH'], 'O': ['HH']})\\
... .get_all_possible_distinct_next_step_count('HOH')
4
>>> Machine({'H': ['HO', 'OH'], 'O': ['HH']})\\
... .get_all_possible_distinct_next_step_count('HOHOHO')
7
"""
return len(self.get_all_possible_distinct_next_steps(start))
def get_all_possible_distinct_next_steps(self, start: str) -> Set[str]:
"""
>>> sorted(Machine({'H': ['HO', 'OH'], 'O': ['HH']})
... .get_all_possible_distinct_next_steps('HOH'))
['HHHH', 'HOHO', 'HOOH', 'OHOH']
"""
return set(self.get_all_possible_next_steps(start))
def get_all_possible_next_steps(self, start: str) -> Iterable[str]:
"""
>>> list(Machine({'H': ['HO', 'OH'], 'O': ['HH']})
... .get_all_possible_next_steps('HOH'))
['HOOH', 'HOHO', 'OHOH', 'HOOH', 'HHHH']
"""
for match, results in self.substitutions.items():
for result in results:
yield from self.get_next_steps_for_replacement(
start, match, result)
def get_next_steps_for_replacement(self, start: str, match: str,
result: str) -> Iterable[str]:
"""
>>> list(Machine().get_next_steps_for_replacement('HOH', 'H', 'HO'))
['HOOH', 'HOHO']
"""
prefix = ""
remaining = start
while len(remaining) >= len(match):
parts = remaining.split(match, 1)
if len(parts) < 2:
return
extra_prefix, suffix = parts
yield f"{prefix}{extra_prefix}{result}{suffix}"
prefix = f"{prefix}{extra_prefix}{match}"
remaining = start[len(prefix):]
Challenge.main()
challenge = Challenge()