1

noi2001_food_chain_dsu.py

@@ -0,0 +1,77 @@
+import dataclasses
+import sys
+
+
+class Dsu:
+    @dataclasses.dataclass
+    class Node:
+        parent: int
+        tree_size: int = 1
+        weight: int = 0
+
+    def __init__(self, n: int):
+        self.nodes: list["Dsu.Node"] = [self.Node(i) for i in range(n)]
+
+    def find(self, x: int) -> int:
+        if self.nodes[x].parent == x:
+            return x
+        parent = self.find(self.nodes[x].parent)
+        self.nodes[x].weight \
+            = (self.nodes[self.nodes[x].parent].weight + self.nodes[x].weight + 12) % 3
+        self.nodes[x].parent = parent
+        return parent
+
+    def union(self, x: int, y: int, weight_offset: int):
+        """
+        x 的父亲为 y；x 的权重为 y 的权重 + weight_offset
+        :param x:
+        :param y:
+        :param weight_offset:
+        :return:
+        """
+        self.find(x)
+        self.find(y)
+        if x == y:
+            return
+        if self.nodes[x].tree_size > self.nodes[y].tree_size:
+            x, y = y, x
+            weight_offset = -weight_offset
+        wx, wy = self.nodes[x].weight, self.nodes[y].weight
+        weight_offset += wy - wx
+        x,y = self.nodes[x].parent, self.nodes[y].parent
+        self.nodes[x].parent = y
+        self.nodes[x].weight = weight_offset % 3
+        self.nodes[y].tree_size += self.nodes[x].tree_size
+
+
+if __name__ == '__main__':
+    animals_amount, instructions_amount = map(int, input().split())
+    dsu = Dsu(animals_amount)
+    lies_amount = 0
+    for _ in range(instructions_amount):
+        instruction_type, animal1, animal2 = map(int, input().split())
+        # print(instruction_type, animal1, animal2, file=sys.stderr)
+        if animal1 > animals_amount or animal2 > animals_amount:
+            print("Lie: No animal", file=sys.stderr)
+            lies_amount += 1
+            continue
+        if animal1 == animal2:
+            if instruction_type == 2:
+                print("Lie: Eats oneself", file=sys.stderr)
+                lies_amount += 1
+            continue
+        animal1 -= 1
+        animal2 -= 1
+        if dsu.find(animal1) == dsu.find(animal2):
+            w1 = dsu.nodes[animal1].weight
+            w2 = dsu.nodes[animal2].weight
+            if instruction_type == 1 and w1 != w2:
+                print("Lie: different type", w1, w2, file=sys.stderr)
+                lies_amount += 1
+            if instruction_type == 2 and (w1 - w2 + 12) % 3 != 1:
+                print("Lie: cannot eat", w1, w2, file=sys.stderr)
+                lies_amount += 1
+        else:
+            dsu.union(animal1, animal2, instruction_type - 1)
+        # print("Success", file=sys.stderr)
+    print(lies_amount)