[PAT-A] 1143 Lowest Common Ancestor

Java

测试点3、4、5超时

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

public class Main {
    static class Node {
        int val;
        Node left, right;

        Node(int val) {
            this.val = val;
        }
    }

    static Node root;

    // 插入节点到BST
    static Node insert(Node node, int val) {
        if (node == null) {
            return new Node(val);
        }
        if (val < node.val) {
            node.left = insert(node.left, val);
        } else {
            node.right = insert(node.right, val);
        }
        return node;
    }

    // 查找节点是否存在
    static boolean find(Node node, int val) {
        if (node == null) {
            return false;
        }
        if (node.val == val) {
            return true;
        }
        return val < node.val ? find(node.left, val) : find(node.right, val);
    }

    // 查找LCA
    static Node findLCA(Node node, int u, int v) {
        if (node == null) {
            return null;
        }
        if (node.val == u || node.val == v) {
            return node;
        }

        // 如果u和v分别在左右子树，则当前节点就是LCA
        if ((u < node.val && v >= node.val) || (u >= node.val && v < node.val)) {
            return node;
        }

        // 否则在左子树或右子树中查找
        if (u < node.val) {
            return findLCA(node.left, u, v);
        } else {
            return findLCA(node.right, u, v);
        }
    }

    // 判断是否为祖先
    static boolean isAncestor(Node node, int ancestor, int descendant) {
        if (node == null) {
            return false;
        }
        if (node.val == descendant) {
            return true;
        }

        if (descendant < node.val) {
            return node.val == ancestor ? find(node.left, descendant) : isAncestor(node.left, ancestor, descendant);
        } else {
            return node.val == ancestor ? find(node.right, descendant) : isAncestor(node.right, ancestor, descendant);
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String[] firstLine = br.readLine().split(" ");
        int M = Integer.parseInt(firstLine[0]);
        int N = Integer.parseInt(firstLine[1]);

        // 读取先序遍历序列并构建BST
        String[] values = br.readLine().split(" ");
        for (int i = 0; i < N; i++) {
            root = insert(root, Integer.parseInt(values[i]));
        }

        // 处理查询
        while (M-- > 0) {
            String[] query = br.readLine().split(" ");
            int u = Integer.parseInt(query[0]);
            int v = Integer.parseInt(query[1]);

            boolean hasU = find(root, u);
            boolean hasV = find(root, v);

            if (!hasU && !hasV) {
                System.out.printf("ERROR: %d and %d are not found.\n", u, v);
            } else if (!hasU) {
                System.out.printf("ERROR: %d is not found.\n", u);
            } else if (!hasV) {
                System.out.printf("ERROR: %d is not found.\n", v);
            } else {
                Node lca = findLCA(root, u, v);
                if (lca.val == u && isAncestor(root, u, v)) {
                    System.out.printf("%d is an ancestor of %d.\n", u, v);
                } else if (lca.val == v && isAncestor(root, v, u)) {
                    System.out.printf("%d is an ancestor of %d.\n", v, u);
                } else {
                    System.out.printf("LCA of %d and %d is %d.\n", u, v, lca.val);
                }
            }
        }
    }
}

C++

#include <iostream>
#include <cstring>
#include <algorithm>
#include <unordered_map>

using namespace std;

const int N = 10010;

int n, Q;
int p[N], depth[N];    // BST: 左<，右>=
unordered_map<int, int> pos;    // pos：原值->离散值
int w[N], pre[N], in[N];        // w[]：离散值->原值

int build(int inL, int inR, int preL, int preR, int d) {
    int root = pre[preL];
    depth[root] = d;

    int k = root;
    if (inL < k) p[build(inL, k - 1, preL + 1, preL + 1 + (k - 1 - inL), d + 1)] = root;
    if (k < inR) p[build(k + 1, inR, preL + 1 + (k - 1 - inL) + 1, preR, d + 1)] = root;

    return root;
}

int main() {
    scanf("%d%d", &Q, &n);

    for (int i = 0; i < n; ++i) {
        scanf("%d", &pre[i]);
        w[i] = pre[i];    // 由先序序列，读入原值存于w[]
    }
    sort(w, w + n);    // 离散化：将原值离散化为[0, n - 1]
    for (int i = 0; i < n; ++i) {
        pos[w[i]] = i;    // 因为是BST，故离散值同时可用作查找子树根的中序位置
        in[i] = i;
    }
    for (int i = 0; i < n; ++i) pre[i] = pos[pre[i]];    // 最后将先序序列离散化

    build(0, n - 1, 0, n - 1, 0);

    while (Q--) {
        int a, b;
        scanf("%d%d", &a, &b);

        if (!pos.count(a) && !pos.count(b)) printf("ERROR: %d and %d are not found.\n", a, b);
        else if (!pos.count(a)) printf("ERROR: %d is not found.\n", a);
        else if (!pos.count(b)) printf("ERROR: %d is not found.\n", b);
        if (!pos.count(a) || !pos.count(b)) continue;

        a = pos[a];
        b = pos[b];
        int a0 = a, b0 = b;
        while (a != b) {
            if (depth[a] > depth[b]) a = p[a];
            else b = p[b];
        }

        if (b == a0) printf("%d is an ancestor of %d.\n", w[a0], w[b0]);
        else if (a == b0) printf("%d is an ancestor of %d.\n", w[b0], w[a0]);
        else printf("LCA of %d and %d is %d.\n", w[a0], w[b0], w[a]);
    }

    return 0;
}