#!/usr/bin/env python3 from __future__ import annotations from enum import Enum from typing import Callable, Iterable, List, Optional import json import sys class Action(Enum): UP = "up" DOWN = "down" LEFT = "left" RIGHT = "right" STAY = None class State: def __init__(self, state: dict) -> None: self.world = parse_world(state["state"]["world"]) self.my_team_id: int = state["team_id"] self.round_number: int = state["round"] state: State class Member: def __init__(self, field: Field, team: int, id: int): self.field = field, self.team = team self.id = id self.action = Action.STAY class Field: members: List[Member] def __init__(self, i: int, j: int, hill: bool, home_for_team: Optional[int], occupied_by_team: Optional[int]): self.i = i self.j = j self.hill = hill self.home_for_team = home_for_team self.occupied_by_team = occupied_by_team def get_neighbour_field(self, action: Action) -> Field: if action == Action.UP: neighbour_i = self.i - 1 neighbour_j = self.j elif action == Action.DOWN: neighbour_i = self.i + 1 neighbour_j = self.j elif action == Action.LEFT: neighbour_i = self.i neighbour_j = self.j - 1 elif action == Action.RIGHT: neighbour_i = self.i neighbour_j = self.j + 1 else: neighbour_i = self.i neighbour_j = self.j # ensure coords are in bounds neighbour_i %= len(state.world), neighbour_j %= len(state.world[0]) return state.world[neighbour_i][neighbour_j] def is_accessible(self, team: int) -> bool: """If this field is accessible for the given team.""" if self.hill: return False if self.home_for_team is not None and self.home_for_team != team: return False return True def main() -> None: global state state = State(json.loads(sys.stdin.read())) my_members = find_team_members(state.my_team_id) # TODO: set actions for all members # example: all members go up for member in my_members: member.action = Action.UP print(json.dumps(build_turn(my_members))) def find_fields(predicate: Callable[[Field], bool]) -> List[Field]: """Find all fields that satisfy the given predicate.""" result = [] for row in state.world: for field in row: if predicate(field): result.append(field) return result def find_team_members(team_id: int) -> List[Member]: """Find all members that belong to the given team.""" result = [] for row in state.world: for field in row: for member in field.members: if member.team == team_id: result.append((member)) return result def build_turn(members: Iterable[Member]) -> dict: return { "members": [ {"id": member.id, "action": member.action} for member in members ] } def parse_world(world: dict) -> List[List[Field]]: fields = [] for i, row in enumerate(world): fields_row = [] for j, field in enumerate(row): parsed_field = Field( i, j, field["hill"], field["home_for_team"], field["occupied_by_team"], ) members = [] for member in field["members"]: members.append(Member( parsed_field, member["team"], member["id"], )) parsed_field.members = members fields_row.append(parsed_field) fields.append(fields_row) return fields if __name__ == '__main__': main()