#!/usr/bin/env python3 import json import sys from typing import Callable, Iterable, List, Optional, Tuple Coords = Tuple[int, int] class State: def __init__(self, state: dict) -> None: self.world = parse_world(state["state"]["map"]) self.my_team_id: int = state["team_id"] self.round_number: int = state["round"] class Member: def __init__(self, team: int, id: int, remaining_rounds: int): self.team = team self.id = id self.remaining_rounds = remaining_rounds self.action: Optional[str] = None class Field: def __init__(self, home_for_team: Optional[int], occupied_by_team: Optional[int], hill: bool, members: List[Member]): self.home_for_team = home_for_team self.occupied_by_team = occupied_by_team self.hill = hill self.members = members state: State = 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, Coords]]: """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 row in world: fields_row = [] for field in row: members = [] for member in field["members"]: members.append(Member( member["team"], member["id"], member["remaining_rounds"] )) fields_row.append(Field( field["home_for_team"], field["occupied_by_team"], field["hill"], members )) fields.append(fields_row) return fields if __name__ == '__main__': main()