aoc2022/02/rps.py

58 lines
1.5 KiB
Python

from enum import IntEnum
import fileinput
from typing import Iterable
SCORE_WIN = 6
SCORE_DRAW = 3
SCORE_LOSS = 0
class Choice(IntEnum):
ROCK = 1
PAPER = 2
SCISSOR = 3
def letter_to_choice(letter: str) -> Choice:
letter = ord(letter)
if letter > 67:
letter += 1
return Choice(letter & 3)
def get_better(i: Choice) -> Choice:
return i + 1 if i != 3 else 1
def rps(opponent: Choice, me: Choice) -> int:
"""Return score for given RPS match"""
if opponent == me: # Draw
return SCORE_DRAW + me.value
elif get_better(opponent) == me:
return SCORE_WIN + me.value
else:
return SCORE_LOSS + me.value
def get_correct(opponent: Choice, goal: str) -> Choice:
if goal == "Y": # draw
return Choice(opponent)
elif goal == "Z": # win
return Choice(get_better(opponent))
else: # lose
return Choice(get_better(get_better(opponent)))
def solve(lines: Iterable[str]) -> int:
total_score = 0
for line in lines:
if len(line) == 4:
total_score += rps(letter_to_choice(line[0]), letter_to_choice(line[2]))
return total_score
def solveb(lines: Iterable[str]) -> int:
total_score = 0
for line in lines:
if len(line) == 4:
total_score += rps(letter_to_choice(line[0]), get_correct(letter_to_choice(line[0]), line[2]))
return total_score
if __name__ == "__main__":
print(solve(fileinput.input()))
print(solveb(fileinput.input()))