Computer >> Hướng Dẫn Máy Tính >  >> Lập Trình >> C++

Truy vấn khoảng cách giữa hai nút trên cây nhị phân O(log N) C++

Cho cây nhị phân N nút và Q truy vấn. Mỗi truy vấn cho 2 nút u, v. Tìm khoảng cách (số cạnh) giữa u và v trong O(log N) mỗi truy vấn.

Công thức

dist(u, v) = depth[u] + depth[v] - 2 * depth[LCA(u, v)]

Cần tìm LCA (Lowest Common Ancestor) nhanh.

Kỹ thuật Euler Tour + Segment Tree (RMQ)

  1. DFS Euler Tour: ghi lại thứ tự duyệt nút vào mảng euler[], độ sâu vào depthEuler[], vị trí xuất hiện đầu tiên của mỗi nút vào first[].
  2. LCA(u, v) = nút có độ sâu nhỏ nhất trong euler[first[u]..first[v]] (RMQ).
  3. Xây Segment Tree trên mảng depthEuler[] để query min O(log N).

Chương trình C++

#include <bits/stdc++.h>
using namespace std;

const int MAXN = 100005;
struct Node { int val; Node *left, *right; Node(int v):val(v),left(nullptr),right(nullptr) {} };

vector euler, depthEuler, first(MAXN, -1), depth(MAXN);
vector nodes;

void dfs(Node* u, int d) {
    if (!u) return;
    depth[u->val] = d;
    first[u->val] = euler.size();
    euler.push_back(u->val);
    depthEuler.push_back(d);
    if (u->left) {
        dfs(u->left, d+1);
        euler.push_back(u->val);
        depthEuler.push_back(d);
    }
    if (u->right) {
        dfs(u->right, d+1);
        euler.push_back(u->val);
        depthEuler.push_back(d);
    }
}

struct SegTree {
    vector> t; // {depth, index_in_euler}
    int n;
    SegTree(const vector& dep) {
        n = dep.size();
        t.resize(4*n);
        build(1, 0, n-1, dep);
    }
    pair combine(pair a, pair b) {
        return (a.first <= b.first) ? a : b;
    }
    void build(int v, int tl, int tr, const vector& dep) {
        if (tl == tr) { t[v] = {dep[tl], tl}; return; }
        int tm = (tl+tr)/2;
        build(v*2, tl, tm, dep);
        build(v*2+1, tm+1, tr, dep);
        t[v] = combine(t[v*2], t[v*2+1]);
    }
    pair query(int v, int tl, int tr, int l, int r) {
        if (l > r) return {INT_MAX, -1};
        if (l == tl && r == tr) return t[v];
        int tm = (tl+tr)/2;
        return combine(query(v*2, tl, tm, l, min(r, tm)),
                       query(v*2+1, tm+1, tr, max(l, tm+1), r));
    }
    int lca(int u, int v) {
        int l = first[u], r = first[v];
        if (l > r) swap(l, r);
        return euler[query(1, 0, n-1, l, r).second];
    }
    int dist(int u, int v) {
        int w = lca(u, v);
        return depth[u] + depth[v] - 2*depth[w];
    }
};

int main() {
    // Xây cây:       1
    //             /   \
    //            2     3
    //           / \   /
    //          4   5 6
    nodes.resize(7);
    for (int i = 1; i <= 6; i++) nodes[i] = new Node(i);
    nodes[1]->left = nodes[2]; nodes[1]->right = nodes[3];
    nodes[2]->left = nodes[4]; nodes[2]->right = nodes[5];
    nodes[3]->left = nodes[6];
    
    dfs(nodes[1], 0);
    SegTree st(depthEuler);
    
    cout << "dist(2,6) = " << st.dist(2,6) << endl; // 3
    cout << "dist(4,1) = " << st.dist(4,1) << endl; // 2
    cout << "dist(5,3) = " << st.dist(5,3) << endl; // 3
    return 0;
}

Kết quả

dist(2,6) = 3
dist(4,1) = 2
dist(5,3) = 3

Lưu ý

Code giả định giá trị nút = chỉ số (1..N). Đối với cây tổng quát, cần map giá trị → chỉ số.