AcWing 240. 食物链(Weighted disjoint set union) Described by Java

A difficult problem is about weighted disjoint set union, to maintain another array distance[] to describe the distance from x to p[x]. And the problem.

Like a tree , the distance from the leaves to root represents the category(0, 1, 2) and the upper can be surplus 3. So the situations: if handle is 1: to check if two number is in one category, if not, res++ if handle is 2: to check the category . Both need to be updated, complexity. We can browse the code:

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintWriter;

public class Main {
          
   
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static PrintWriter pw = new PrintWriter(System.out);
    static final int N = 50010;
    static int p[] = new int[N], d[] = new int[N], n, k;

    public static int find(int x) {
          
   
        if (p[x] != x) {
          
   
            int t = find(p[x]);
            d[x] += d[p[x]];
            p[x] = t;
        }
        return p[x];
    }

    public static void main(String[] rags) throws IOException {
          
   
        String[] s = br.readLine().split(" ");
        n = Integer.parseInt(s[0]);
        k = Integer.parseInt(s[1]);
        for (int i = 1; i <= n; i++) p[i] = i;
        int res = 0;
        while (k-- > 0) {
          
   
            String ss[] = br.readLine().split(" ");
            int t = Integer.parseInt(ss[0]);
            int x = Integer.parseInt(ss[1]);
            int y = Integer.parseInt(ss[2]);
            if (x > n || y > n) res++;
            else {
          
   
                int px = find(x), py = find(y);
                if (t == 1) {
          
   
                    if (px == py && (d[x] - d[y]) % 3 != 0) res++;
                    else {
          
   
                        if (px != py) {
          
   
                            p[px] = py;
                            d[px] = d[y] - d[x];
                        }
                    }
                } else {
          
   
                    if (px == py && (d[x] - 1 - d[y]) % 3 != 0) res++;
                    else {
          
   
                        if (px != py) {
          
   
                            p[px] = py;
                            d[px] = d[y] + 1 - d[x];
                        }
                    }
                }
            }
        }
        pw.println(res);
        pw.flush();
        br.close();
    }
}