LCA(최소 공통 조상)

트리의 LCA를 찾아보자.

브루트포스 O(depth)

int depth[N]; // 트리의 깊이가 저장돼있음
int p[N]; // 부모가 저장돼있음

int lca(int u, int v) {
    // WLOG, u < v
    if (depth[u] > depth[v]) swap(u, v);
    // depth가 같아질 때까지 더 낮은걸 올림.
    while (depth[v] != depth[u]) v = p[v];
    // 그리고 만날 때 까지 올림
    while (u != v) u = p[u], v = p[v];

    return u;
}

깊이에 비례하는 시간이 걸린다.

sparse table O(log N)

int p[N+1][logN+1]; // i의 2^j위의 조상
int depth[N+1];

void init() {
    // 루트의 부모를 0으로따로 나타내기위해 노드 번호 1~N으로 세팅
    // p[i][0]은 본인의 직계 부모로 초기화돼있다고 가정
    // p[root][0] = 0;
    for (int j = 1; (1 << j) < N; ++j) {
        for (int i = 1; i <= N; ++i) {
            p[i][j] = p[p[i][j - 1]][j - 1];
        }
    }
}

int lca(int u, int v) {
    // WLOG, depth[u] < depth[v]
    if (depth[u] > depth[v]) swap(u, v);

    int k = 0; 
    while ((1 << (k+1)) <= depth[v]) k++;

    // v를 u까지 올리기
    for (int i = k; i >= 0; --i) {
        if (depth[v] - (1 << i) >= depth[u])
            v = p[v][i];
    }
    if (u == v) return u;

    // 같이 올리기(직전까지만)
    for (int i = k; i >= 0; --i) {
        if (p[u][i] != 0 && p[u][i] != p[v][i])
            u = p[u][i], v = p[v][i];
    }
    return p[u][0];
}

희소배열의 성질을 이용해 log N만에 lca 쿼리를 해결할 수 있다.
같이 올릴 때 직전까지만 올리는 이유를 이해한다면 희소배열을 완벽히 이해한 것?

Sparse Table 2 O(log N)

vector<int> adj[N+1]; // 그래프
int p[N+1][logN +1];
int in[N+1]; int out[N+1]; // depth대신 순서
int order = 0;
// 1~N, 루트가 1이면 (1,1)으로 시작
void dfs(int here, int parent) {
	in[here] = ++order;
    
    //dfs특성상, 부모의 p정보는 이미다 구현되어 있다.
    p[here][0] = parent;
    for (int j=1; j< log N +1; ++j) 
    	P[here][j] = p[p[here][j-1]][j-1];
    for (int next : adj[here])
    	if (next != parent) dfs(next, here);
    
    out[here] = ++order;
}

inline bool isAncestor(int u, int v) {
	return in[u]<=in[v] && out[v]<=out[u];
}

int lca(int u, int v) {
	if (isAncestor(u, v)) return u;
    if (isAncestor(v, u)) return v;
    
    // 직전까지 올리기
    for (int k = logN; k>=0; --k) {
    	if ( !isAncestor(p[u][k], v))
        	u = p[u][k];
    }
    return p[u][0];
}

지금까지는 주어진 그래프 정보를 이용해 bfs로 트리를 만들었다고 가정했다. 하지만 dfs로 트리를 만들게 되면 여러 성질을 이용할 수 있다. depth대신 dfs의 순회 순서를 이용해서 u, v간의 조상관계를 파악할 수 있다. 그리고 dfs는 무조건 모든 조상을 이전에 거쳐왔다는 성질을 이용해서 트리를 구성하면서 희소배열을 동시에 초기화 할 수 있다.

시간복잡도는 이전의 희소배열과 같지만 bfs + 희소배열 bottom up -> dfs로 초기화가 빨라지고, isAncestor를 통해 lca의 평균 시간복잡도를 줄일 수 있다. 하지만 막상 돌려보면 이전 코드랑 비슷하게 걸린다..?

bfs로 트리를 만드느냐 dfs로 트리를 만드느냐의 차이다.

세그먼트 트리를 이용한 LCA

void dfs(int here, int d, int parent) {
    first[here] = order.size();
    order.push_back(here);
    depth.push_back(d);

    for (int next : adj[here]) {
        if (next == parent) continue;
        dfs(next, d + 1, here);
        order.push_back(here);
        depth.push_back(d);
    }
}

int lca(SegTree& segtree, int u, int v) {
    int a = first[u], b = first[v];
    if (b < a) swap(a, b);

    // a~b구간에서 depth가 가장 낮은 노드 리턴
    int idx =  segtree.query(a, b);
    return order[idx];
}

struct SegTree {
    int n;
    vi tree;

    //Constructor
    SegTree(const vector<int>& arr) {
        n = arr.size();
        tree.resize(n * 4); // 1<<(log2(n)+1)로 해도된다
        init(arr, 0, n - 1, 1);
    }

    // 전처리 O(NlogN)
    int init(const vector<int>& arr, int lo, int hi, int node) {
        if (lo == hi)
            return tree[node] = lo;

        int mid = (lo + hi) / 2;
        int leftVal = init(arr, lo, mid, node * 2);
        int rightVal = init(arr, mid + 1, hi, node * 2 + 1);

        if (depth[leftVal] <= depth[rightVal])
            return tree[node] = leftVal;
        else
            return tree[node] = rightVal;
    }

    // Query O(logN)
    int query(int lo, int hi) {
        return query(lo, hi, 1, 0, n - 1);
    }
    int query(int lo, int hi, int node, int nodelo, int nodehi) {
        if (nodehi < lo || hi < nodelo) return INF;
        if (lo <= nodelo && nodehi <= hi) {
            return tree[node];
        }

        int mid = (nodelo + nodehi) / 2;
        int leftVal = query(lo, hi, node * 2, nodelo, mid);
        int rightVal = query(lo, hi, node * 2 + 1, mid + 1, nodehi);
        if (leftVal == INF) return rightVal;
        if (rightVal == INF) return leftVal;
        if (depth[leftVal] <= depth[rightVal])
            return leftVal;
        else
            return rightVal;
    }
};

dfs를 이용해서 순회 순서대로 노드번호와 depth를 기록해둔다. 만든 order벡터에서 구간 중에 depth가 가장 낮은 값을 구하는 segtree를 만들어서 lca를 해결할 수 있다. 아이디어만 흡수하면 되고, lca 구하는데는 sparse table 사용하자.

Reference

• 백준님 강의