ABC087 D — People on a Line

1# -*- coding: utf-8 -*-2 3 4class WeightedUnionFind:5    """Represents a data structure that tracks a set of elements partitioned6       into a number of disjoint (non-overlapping) subsets.7    Landau notation: O(α(n)), where α(n) is the inverse Ackermann function.8    See:9    https://www.youtube.com/watch?v=zV3Ul2pA2Fw10    https://en.wikipedia.org/wiki/Disjoint-set_data_structure11    https://atcoder.jp/contests/abc120/submissions/444494212    https://qiita.com/drken/items/cce6fc5c579051e64fab13    """14 15    def __init__(self, number_count: int):16        """17        Args:18            number_count: The size of elements (greater than 2).19        """20        self.parent_numbers = [i for i in range(number_count)]21        self.rank = [0 for _ in range(number_count)]22        self.diff_weight = [0 for _ in range(number_count)]23 24    def find_root(self, number: int) -> int:25        """Follows the chain of parent pointers from number up the tree until26           it reaches a root element, whose parent is itself.27        Args:28            number: The trees id (0-index).29        Returns:30            The index of a root element.31        """32        if self.parent_numbers[number] == number:33            return number34        else:35            parent_number = self.parent_numbers[number]36            root = self.find_root(parent_number)37            self.diff_weight[number] += self.diff_weight[parent_number]38            self.parent_numbers[number] = root39            return root40 41    def calc_weight(self, number: int) -> int:42        self.find_root(number)43        return self.diff_weight[number]44 45    def is_same_group(self, number_x: int, number_y: int) -> bool:46        """Represents the roots of tree number_x and number_y are in the same47           group.48        Args:49            number_x: The trees x (0-index).50            number_y: The trees y (0-index).51        """52        return self.find_root(number_x) == self.find_root(number_y)53 54    def merge_if_needs(self, number_x: int, number_y: int, weight: int) -> bool:55        """Uses find_root to determine the roots of the tree number_x and56           number_y belong to. If the roots are distinct, the trees are combined57           by attaching the roots of one to the root of the other.58        Args:59            number_x: The trees x (0-index).60            number_y: The trees y (0-index).61        """62        # Correct the difference between the weight of root and number_x, number_y63        weight += self.calc_weight(number_x)64        weight -= self.calc_weight(number_y)65        root_x, root_y = self.find_root(number_x), self.find_root(number_y)66 67        if root_x == root_y:68            return False69 70        if self.rank[root_x] < self.rank[root_y]:71            root_x, root_y = root_y, root_x72            weight = -weight73        if self.rank[root_x] == self.rank[root_y]:74            self.rank[root_x] += 175 76        self.parent_numbers[root_y] = root_x77        self.diff_weight[root_y] = weight78        return True79 80    def calc_cost(self, from_x: int, to_y: int) -> int:81        return self.calc_weight(to_y) - self.calc_weight(from_x)82 83 84def main():85    n, m = map(int, input().split())86    wuf = WeightedUnionFind(n)87 88    for i in range(m):89        li, ri, di = map(int, input().split())90        li -= 191        ri -= 192 93        if wuf.is_same_group(li, ri):94            diff = wuf.calc_cost(li, ri)95 96            if diff != di:97                print("No")98                exit()99        else:100            wuf.merge_if_needs(li, ri, di)101 102    print("Yes")103 104 105if __name__ == "__main__":106    main()107