#!/usr/bin/env python3 from __future__ import annotations from typing import Callable, Iterable, Optional, List, Set import collections import enum import json import sys class Action(enum.Enum): UP = enum.auto() DOWN = enum.auto() LEFT = enum.auto() RIGHT = enum.auto() STAY = enum.auto() _inversions = { UP: DOWN, DOWN: UP, LEFT: RIGHT, RIGHT: LEFT, STAY: STAY, } def invert(self) -> Action: return self._inversions[self] def __str__(self) -> str: return self.name.lower() 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) -> None: 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], protected_for_team: Optional[int],) -> None: self.i = i self.j = j self.hill = hill self.home_for_team = home_for_team self.occupied_by_team = occupied_by_team self.protected_for_team = protected_for_team def __eq__(self, other: Field) -> bool: return (self.i, self.j) == (other.i, other.j) def __hash__(self) -> int: return hash((self.i, self.j)) 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.protected_for_team is not None and self.protected_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 pathfind(member: Member, goal: Field) -> Action: """Find the shortest path from member position to goal. Returns the first action to take. If there is no path, returns Action.STAY. """ # BFS from goal to member # this direction makes it easier to realize that no path exists explored: Set[Field] = set() queue = collections.deque([goal]) while queue: field = queue.popleft() for action in Action: neighbour = field.get_neighbour_field(action) if (neighbour in explored or not neighbour.is_accessible(member.team)): continue if neighbour == member.field: return action.invert() queue.append(neighbour) explored.add(neighbour) return Action.STAY def build_turn(members: Iterable[Member]) -> dict: return { "members": [ {"id": member.id, "action": str(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"], field["protected_for_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()