1import math
  2
  3
  4class Candy:
  5    """Represent a Halloween in Pripyat candy."""
  6
  7    def __init__(self, mass, uranium):
  8        """Initializator."""
  9        self._mass = mass
 10        self._uranium = uranium
 11    
 12    def get_uranium_quantity(self):
 13        """Get absolute content of uranium."""
 14        return self._mass * self._uranium
 15    
 16    def get_mass(self):
 17        """Get the mass of the candy."""
 18        return self._mass
 19
 20
 21class Person:
 22    """Represent a Halloween in Pripyat person."""
 23
 24    def __init__(self, position):
 25        """Initializator."""
 26        self._position = position
 27
 28    def get_position(self):
 29        """Get the current position of a person."""
 30        return self._position
 31
 32    def set_position(self, position):
 33        """Set the current position of a person."""
 34        self._position = position
 35
 36
 37class Kid(Person):
 38    """Represent a Halloween in Pripyat kid."""
 39
 40    CRITICAL_URANIUM = 20
 41
 42    def __init__(self, position, initiative):
 43        """Initializator."""
 44        super().__init__(position)
 45        self._candies = set()
 46        self._initiative = initiative
 47
 48    def get_initiative(self):
 49        """Get the initiative value of the kid."""
 50        return self._initiative
 51
 52    def add_candy(self, candy):
 53        """Add a new candy to the set of owned candies."""
 54        self._candies.add(candy)
 55    
 56    def is_critical(self):
 57        """Check if the uranium mass in owned candies is critical."""
 58        total = sum([candy.get_uranium_quantity() for candy in self._candies])
 59        return total > self.CRITICAL_URANIUM
 60
 61
 62class Host(Person):
 63    """Represent a Halloween in Pripyat host."""
 64
 65    def __init__(self, position, candies):
 66        """Initializator."""
 67        super().__init__(position)
 68        self._candies = {Candy(*candy) for candy in candies}
 69
 70    def remove_candy(self, fun):
 71        """Remove a candy from the set of owned candies."""
 72        if not self._candies:
 73            return None
 74        candy = fun(self._candies)
 75        self._candies.remove(candy)
 76        return candy
 77
 78
 79class Visit:
 80    """A single visit between a kid and a host."""
 81
 82    def __init__(self, kid, host, turn_index):
 83        """Initializator."""
 84        self.kid = kid
 85        self.host = host
 86        self.turn_index = turn_index
 87
 88
 89class FluxCapacitor:
 90    """Represent a Halloween in Pripyat simulator."""
 91
 92    def __init__(self, participants):
 93        """Initializator."""
 94        self._kids = list(filter(lambda x: type(x) is Kid, participants))
 95        self._hosts = list(filter(lambda x: type(x) is Host, participants))
 96        self._visits = self._init_visits()
 97
 98    @staticmethod
 99    def _distance(kid, host):
100        """Calculate dinstance between a kid and a host."""
101        return math.dist(kid.get_position(), host.get_position())
102
103    @staticmethod
104    def _select_candy(candies):
105        """Return a function to select the candy with highest mass from a set of candies."""
106        return max(candies, key=lambda candy: candy.get_mass())
107    
108    @classmethod
109    def _get_closest_host_of_kid(cls, kid, hosts):
110        """Get the host that is closest to a particular kid."""
111        return min(hosts, key=lambda host: (cls._distance(kid, host), host.get_position()))
112    
113    def _init_visits(self):
114        """Calculate all visits for the night."""
115        visits = []
116        for kid in self._kids:
117            hosts = self._hosts[:]
118            turn_index = 0
119            while hosts:
120                host = self._get_closest_host_of_kid(kid, hosts)
121                hosts.remove(host)
122                visits.append(Visit(kid, host, turn_index))
123                turn_index += 1
124                kid.set_position(host.get_position())
125        return visits
126
127    def _pop_visits(self):
128        """Remove next visits from all-visits set and return them."""
129        popped_visits = []
130        next_turn = min(self._visits, key=lambda visit: visit.turn_index).turn_index
131        for visit in self._visits:
132            if visit.turn_index == next_turn:
133                popped_visits.append(visit)
134        self._visits = [visit for visit in self._visits if visit not in popped_visits]
135        return popped_visits
136
137    def _make_a_visit(self):
138        """Make a single visit for the night and return exchanges for them."""
139        if not self._visits:
140            return None
141        visits = self._pop_visits()
142        exchanges = {}
143        for visit in visits:
144            exchanges.setdefault(visit.host, []).append(visit.kid)
145        return {host: sorted(kids, key=lambda kid: kid.get_initiative(), reverse=True) for host, kids in exchanges.items()}
146
147    def _exchange_candy(self, exchanges):
148        """Make a candy exchange."""
149        deads = set()
150        for host, kids in exchanges.items():
151            for kid in kids:
152                candy = host.remove_candy(self._select_candy)
153                if candy:
154                    kid.add_candy(candy)
155                    if kid.is_critical():
156                        deads.add(kid)
157        return deads or None
158
159    def get_victim(self):
160        """Get the first victim of the night."""
161        while True:
162            if not (exchanges:=self._make_a_visit()):
163                return None
164            if victim:=self._exchange_candy(exchanges):
165                return victim