样例2输出错误求调
查看原帖
样例2输出错误求调
1152063
SuperYC233楼主2024/6/6 11:56

用的倍增lca

import java.io.*;
import java.util.*;

public class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static PrintWriter pw = new PrintWriter(new BufferedOutputStream(System.out));
    static int K = 20;
    static int N = (int) (1e5 + 10);
    static int M = N * 2;
    static int n, m;
    static int[] c = new int[N];
    static int[] h = new int[N];
    static int[] e = new int[M];
    static int[] ne = new int[M];
    static int[][] fa = new int[N][K + 1];
    static int[] depth = new int[N];
    static int idx;

    public static void main(String[] args) throws IOException {
        String[] strs = br.readLine().split(" ");
        n = Integer.parseInt(strs[0]);
        m = Integer.parseInt(strs[1]);
        strs = br.readLine().split(" ");
        for (int i = 1; i <= n; i++) {
            c[i] = Integer.parseInt(strs[i - 1]);
        }

        Arrays.fill(h, -1);
        for (int i = 1; i < n ; i++) {
            strs = br.readLine().split(" ");
            int a = Integer.parseInt(strs[0]);
            int b = Integer.parseInt(strs[1]);
            add(a, b);
            add(b, a);
        }

        bfs();

        for (int i = 1; i <= m; i++) {
            strs = br.readLine().split(" ");
            int a = Integer.parseInt(strs[0]);
            int b = Integer.parseInt(strs[1]);
            pw.println(lca(a, b));
        }
        pw.flush();
    }

    public static int lca(int a, int b) {
        HashSet<Integer> set = new HashSet<>();
        if (depth[a] < depth[b]) {
            int tmp = b;
            b = a;
            a = tmp;
        }
        set.add(c[a]);
        set.add(c[b]);

        // 先跳到同一高度
        for (int k = K; k >= 0; k--) {
            if (depth[fa[a][k]] >= depth[b]) {
                a = fa[a][k];
                set.add(c[a]);
            }
        }
        if (a == b) return set.size();

        for (int k = K; k >= 0; k--) {
            if (fa[a][k] != fa[b][k]) {
                a = fa[a][k];
                b = fa[b][k];
                set.add(c[a]);
                set.add(c[b]);
            }
        }

        set.add(c[fa[a][0]]);
        return set.size();
    }

    public static void bfs() {
        Queue<Integer> queue = new LinkedList<>();
        Arrays.fill(depth, 0x3f3f3f3f);
        depth[0] = 0;
        depth[1] = 1;
        queue.add(1);
        while (!queue.isEmpty()) {
            int t = queue.poll();
            for (int i = h[t]; i != -1; i = ne[i]) {
                int j = e[i];
                if (depth[j] > depth[t] + 1) {
                    depth[j] = depth[t] + 1;
                    queue.add(j);
                    fa[j][0] = t;
                    for (int k = 1; k <= K; k++) {
                        fa[j][k] = fa[fa[j][k - 1]][k - 1];
                    }
                }
            }
        }
    }

    public static void add(int a, int b) {
        e[idx] = b;
        ne[idx] = h[a];
        h[a] = idx++;
    }
}
2024/6/6 11:56
加载中...