101 lines
2.3 KiB
Python
101 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os.path
|
|
import re
|
|
|
|
import pytest
|
|
|
|
import support
|
|
|
|
INPUT_TXT = os.path.join(os.path.dirname(__file__), 'input.txt')
|
|
|
|
|
|
def compute(s: str):
|
|
_stacks = []
|
|
for _ in range(9):
|
|
_stacks.append([])
|
|
stopped_at_idx = None
|
|
for i, line in enumerate(s.split('\n')):
|
|
if not line.strip().startswith("["):
|
|
stopped_at_idx = i
|
|
break
|
|
for idx, char in enumerate(line):
|
|
if char in "[]":
|
|
continue
|
|
elif char != " ":
|
|
if idx == 1:
|
|
stack_idx = 0
|
|
else:
|
|
stack_idx = ((idx - 1) // 4)
|
|
_stacks[stack_idx].append(char)
|
|
stacks = []
|
|
for st in _stacks:
|
|
stacks.append(list(reversed(st)))
|
|
|
|
if stopped_at_idx is None:
|
|
raise Exception
|
|
for line in s.split('\n'):
|
|
if not line.strip().startswith("move"):
|
|
print(line)
|
|
continue
|
|
match = re.match(r".* ([0-9]+).*([0-9]+).*([0-9]+)", line)
|
|
quantity, fr, t = match[1], match[2], match[3] # type: ignore
|
|
print(quantity, fr, t)
|
|
fr = int(fr) - 1
|
|
t = int(t) - 1
|
|
q = int(quantity)
|
|
print(stacks[fr], stacks[t])
|
|
|
|
if len(stacks[fr]) < q:
|
|
to_move = stacks[fr]
|
|
stacks[t] += to_move
|
|
stacks[fr] = []
|
|
else:
|
|
to_move = stacks[fr][-q:]
|
|
stacks[t] += to_move
|
|
stacks[fr] = stacks[fr][:-q]
|
|
|
|
print(stacks[fr], stacks[t])
|
|
answer = ''
|
|
for stack in stacks:
|
|
if stack:
|
|
answer += stack[-1]
|
|
else:
|
|
print('appending a blank')
|
|
answer += " "
|
|
print('answer:', answer)
|
|
return answer
|
|
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
('input_s', 'expected'),
|
|
((""" [D]
|
|
[N] [C]
|
|
[Z] [M] [P]
|
|
1 2 3
|
|
|
|
move 1 from 2 to 1
|
|
move 3 from 1 to 3
|
|
move 2 from 2 to 1
|
|
move 1 from 1 to 2""", "CMZ"),
|
|
))
|
|
def test(input_s: str, expected: int) -> None:
|
|
assert compute(input_s) == expected
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('data_file', nargs='?', default=INPUT_TXT)
|
|
args = parser.parse_args()
|
|
|
|
with open(args.data_file) as f, support.timing():
|
|
print(compute(f.read()))
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == '__main__':
|
|
raise SystemExit(main())
|