Editorial for Pokemon

Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

Approach

Model each pokemon's type as an integer mod $t$ ~t~. The statement " $a$ ~a~ beats $b$ ~b~" means $\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 $t$ ~t~). When uniting on " $a$ ~a~ beats $b$ ~b~":

If $a$ ~a~ and $b$ ~b~ 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)

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")

Comments

There are no comments at the moment.