#!/usr/bin/env python3 import json import sys from typing import Callable, Iterable, List, Optional, Tuple 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 = None class Member: def __init__(self, field: 'Field', team: int, id: int): self.field = field, self.team = team self.id = id self.action: Optional[str] = None class Field: 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 self.members: List[Member] = None 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 = "up" print(json.dumps(build_turn(my_members))) def find_fields( predicate: Callable[[Field], bool] ) -> List[Tuple[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[Tuple[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 ) fields.append(parsed_field) members = [] for member in field["members"]: members.append(Member( parsed_field, member["team"], member["id"], )) parsed_field.members = members fields.append(fields_row) return fields if __name__ == '__main__': main()