Editorial for Pokemon


Approach

Model each pokemon's type as an integer mod tt. The statement "aa beats bb" means type(a)≡type(b)+1(modt)\mathrm{type}(a) \equiv \mathrm{type}(b) + 1 \pmod t.

Use a weighted DSU where each node stores its type offset relative to its parent (mod tt). When uniting on "aa beats bb":

  • If aa and bb share a root, check whether their stored offsets already satisfy the relation.
  • Otherwise link roots and set the new offset so the relation holds.

Inconsistent statements are rejected (No) and ignored afterwards. Runtime is nearly O(n+q)O(n + q).

Solution (Python)

Code 1
n,t,q = (int(x) for x in input().split())

class DisjointSet():

    def __init__(self,size : int) -> None:
        self.parent = [i for i in range(size)]
        self.size = [1 for _ in range(size)]
        self.rel_to_parent = [0 for _ in range(size)]

    def find_set(self,v : int) -> tuple[int,int]:

        if (self.parent[v] == v): return (v,0)

        parent,value = self.find_set(self.parent[v])
        self.parent[v] = parent
        self.rel_to_parent[v] = (value + self.rel_to_parent[v]) % t

        return (self.parent[v],self.rel_to_parent[v])
    
    def union(self, v : int, u : int) -> None:
        """
        if v beats u, then v_val = u_val + 1 should end up being true.
        """
        v_par,v_val = self.find_set(v)
        u_par,u_val = self.find_set(u)

        if v_par == u_par:
            return v_val == (u_val + 1) % t

        if (self.size[v_par] > self.size[u_par] ):
            # Attaching u_par to v_par
            self.parent[u_par] = v_par
            self.rel_to_parent[u_par] = (v_val - u_val - 1 ) % t

            self.size[v_par] += self.size[u_par]
        else:
            # Attaching v_par to u_par
            self.parent[v_par] = u_par
            self.rel_to_parent[v_par] = (u_val + 1 - v_val) % t

            self.size[u_par] += self.size[v_par]

        return True

dsu = DisjointSet(n+1)

for _ in range(q):
    a,b = (int(x) for x in input().split())
    print("Yes" if dsu.union(a,b) else "No")

Comments0


No comments yet

Be the first to comment.

New comment


Log in to join the discussion.