[LeetCode] Binary Lifiting를 사용한 최소공통조상 최적화

최근 LeetCode 문제를 풀며 최소공통조상 비슷한 문제가 나와 쉽게 풀었다고 생각했다.

 

https://leetcode.com/problems/number-of-ways-to-assign-edge-weights-ii/description/?envType=daily-question&envId=2026-06-12

Number of Ways to Assign Edge Weights II - LeetCode

 

문제를 보면

1. n개의 노드가 트리 구조로 주어진다. 이때 루트는 1이다.

2. 노드들을 연결하는 edges가 2차원 배열로 주어진다.

3. 모든 간선의 가중치(무게)가 0이지만, 각 간선에 1 또는 2를 부여할 수 있다.

4. queries 배열을 통해 2개의 노드가 주어지는데 a 노드에서 b 노드로 가는 간선중에 1, 2를 부여한다.

5. 간선의 가중치의 합이 홀수가 되는 모든 경우의 수를 구하여 10^9+7로 나누어야한다.

 

간선에 부여할 수 있는 숫자는 1, 2 총 2가지이다.

 

그렇다면 나올 수 있는 경우의 수를 살펴보자

 

간선이 한개인 경우 2개인 경우 3개인 경우를 보면

 

한개인 경우 : 1, 2 (2가지)

두개인 경우 : (1, 1), (1, 2), (2, 1), (2, 2) 4가지

세개인 경우 : (1, 1, 1), (1, 1, 2), (1, 2, 1), (2, 1, 1), (1, 2, 2), (2, 1, 2), (2, 2, 1), (2, 2, 2) 8가지

 

어차피 각각의 간선에 2가지가 올 수 있으니까 간선의 개수를 n이라고 하면 2^n 아닌가요? 할 수 있다.

 

물론 맞는 말이다. 하지만 일일히 이렇게 구한 이유가 있다.

 

저기서 공통적으로 찾을 수 있는 것은 한개인 경우, 두개인 경우, 세개인 경우 모두 다 합이 홀수, 짝수 반반씩이다

즉 두개인 경우 4가지 중 간선의 합이 홀수인 경우는 2가지 짝수인 경우는 2가지로 나누어 떨어진다.

 

해당 규칙을 찾으면 이제 최소공통조상을 찾아서 간선의 개수만 찾으면 답을 쉽게 구할 수 있다.

 

최소 공통 조상(LCA)

최소 공통 조상을 구하기 위해 작성한 코드이다.

class Solution {
typedef long long ll;
public:
    vector<int> assignEdgeWeights(vector<vector<int>>& edges, vector<vector<int>>& queries) {
        vector<vector<int>> v(edges.size()+2);
        vector<int> answer(queries.size(), 0);
        vector<int> height(edges.size()+2);
        vector<int> parent(edges.size()+2);

        for(auto &iter : edges) {
            int a = iter[0];
            int b = iter[1];

            v[a].push_back(b);
            v[b].push_back(a);
        }

        dfs(1, 0, v, parent, 0, height);

        for(int i = 0; i < queries.size(); i++) {
            int a = queries[i][0];
            int b = queries[i][1];
            int len = 0;

            while(a != b) {
                if(height[a] > height[b]) {
                    len++;
                    a = parent[a];
                }
                else if(height[a] < height[b]) {
                    len++;
                    b = parent[b];
                }
                else {
                    a = parent[a];
                    b = parent[b];
                    len += 2;
                }
            }

            if(len == 0) {
                answer[i] = 0;
            }
            else answer[i] = Pow(len-1);
        }

        return answer;
    }
};

 

dfs 함수는 저기서 중요한게 아니니까 제거했다. 여기서 공통 조상을 구하기 위해 depth를 맞춰주고 노드의 부모를 계속해서 확인하며 거슬러 올라간다. 쿼리의 크기는 10^5, 노드의 개수도 10^5으로 시간복잡도를 계산하면 10^10이 될 수 있다.

 

해당 문제를 해결하기 위해 Binary Lifting 기법을 사용해서 log n으로 찾을 수 있게 구현해보았다.

 

Binary Lifting

기존 1씩 거슬러 올라가기에는 시간이 오래걸리기에 2의 제곱만큼 거슬러 올라가는 방식이다.

 

 

10 노드와 12 노드를 보면 1칸 올라가는 부모는 각각 6과 7이다. 그렇다면 2칸 올라갔을때 부모는 공통적으로 4이다.

dp 테이블을 2의 거듭제곱만큼 올라가는 횟수로 지정하고 열은 노드로 설정해서 테이블에 기록해준다.

 

dp[r][0] = r이 0인 경우 해당 노드에 대한 한칸올라갔을때의 부모이며 미리 기록해주었으며

 

2칸올라갔을때 부모를 구할때는 1칸 올라갔을때의 노드의 부모를 찾고 해당 노드의 이전 제곱근 즉 한칸 올라갔을때의 부모를 찾으면 그것이 2칸올라갔을때의 부모이며

 

해당 내용을 점화식으로 나타내면 dp[i][j] = dp[i-1][dp[i-1][j]] 로 표현 가능하다.

 

        vector<vector<int>> dp(log, vector<int>(edges.size()+2, 0));

        for(int i = 0; i < edges.size()+2; i++) {
            dp[0][i] = parent[i];
        }

        for(int i = 1; i < log; i++) {
            for(int j = 0; j < edges.size()+2; j++) {
                dp[i][j] = dp[i-1][dp[i-1][j]];
            }
       }
}

 

이제 공통 조상을 찾기 위해서 거슬러 올라가야 한다.

 

먼저 LCA와 같이 높이가 다른 경우 공통의 조상을 찾을 수 없기에 높이를 맞춰줘야한다.

 

        auto lca = [&](int u, int w) {
            if(u == w) return 0;
            if(height[u] < height[w]) swap(u, w);

 

u, w에 대한 노드가 있을때 해당 노드가 같은 경우 간선의 개수는 0이기에 리턴해주었고 필자는 u를 가장 큰 depth로 두었다. w 노드의 depth가 큰 경우 u와 w를 변경해주었다.

 

int diff = height[u] - height[w];

for(int k = 0; k < log; k++) {
    if(diff & (1 << k)) u = dp[k][u];
}

 

diff 변수가 의미하는 것은 높이의 차이이다. 즉 높이를 맞춰주기위해 변수에 담아 주었고

k가 0부터 log (로그가 의미하는 것은 log2 트리의 높이) 만큼 반복하며 높이 차이만큼 거슬러 올라가주었다.

 

if(u == w) return diff;

int value = diff;
            
for(int k = log - 1; k >= 0; k--) {
    if(dp[k][u] != dp[k][w]) {
        u = dp[k][u], w = dp[k][w];

        value += (1 << k) * 2;
        }
    }

	return value + 2;
};

 

이때 u와 w가 같은경우 공통 조상이며 간선의 개수를 리턴해주었다.

 

이후 value라는 변수에 간선의 개수를 기록해주었고

 

여기서 가장 중요한 부분은 log - 1 가장 큰 보폭부터 탐색해서 노드가 다른경우(달라진다는 것은 가지가 두갈래로 나뉜다는 것)에 각각의 노드를 업데이트 해준다.

 

가장 큰 보폭으로 탐색하는 이유는 2의 거듭제곱들의 합으로 쪼개는 것인데 트리의 높이만큼 넘지 않는 선에서 빼줘야지 2진수로 표현이 가능하기 때문이다.

 

이후 간선 값을 2의 거듭제곱으로 계산 후 2를 곱해주어 누적해주었으며 마지막에는 한칸만 올라가면 공통 조상 노드이기에 2를 더해주었다.

 

전체코드

 

class Solution {
typedef long long ll;
private:
    int MOD = 1e9 + 7;
    int max_height = 0;
    void dfs(int node, int par, vector<vector<int>> &v, vector<int> &parent, int hei, vector<int> &height) {
        height[node] = hei;
        parent[node] = par;
        max_height = max(max_height, hei);
        for(auto &iter : v[node]) {
            if(iter != par) {
                dfs(iter, node, v, parent, hei+1, height);
            }
        }
    }

    int Pow(int a) {
        ll val = 2;
        ll result = 1;

        while(a != 0) {
            if(a & 1) {
                result = (result * val) % MOD;
            }

            a >>= 1;
            val = (val * val) % MOD;
        }

        return result % MOD;
    }
public:
    vector<int> assignEdgeWeights(vector<vector<int>>& edges, vector<vector<int>>& queries) {
        vector<vector<int>> v(edges.size()+2);
        vector<int> answer(queries.size(), 0);
        vector<int> height(edges.size()+2);
        vector<int> parent(edges.size()+2);

        for(auto &iter : edges) {
            int a = iter[0];
            int b = iter[1];

            v[a].push_back(b);
            v[b].push_back(a);
        }

        dfs(1, 0, v, parent, 0, height);

        int log = 0;

        while(max_height != 0) {
            log++;
            max_height >>= 1;
        }

        vector<vector<int>> dp(log, vector<int>(edges.size()+2, 0));

        for(int i = 0; i < edges.size()+2; i++) {
            dp[0][i] = parent[i];
        }

        for(int i = 1; i < log; i++) {
            for(int j = 0; j < edges.size()+2; j++) {
                dp[i][j] = dp[i-1][dp[i-1][j]];
            }
        }

        auto lca = [&](int u, int w) {
            if(u == w) return 0;
            if(height[u] < height[w]) swap(u, w);
            int diff = height[u] - height[w];

            for(int k = 0; k < log; k++) {
                if(diff & (1 << k)) u = dp[k][u];
            }

            if(u == w) return diff;

            int value = diff;
            
            for(int k = log - 1; k >= 0; k--) {
                if(dp[k][u] != dp[k][w]) {
                    u = dp[k][u], w = dp[k][w];

                    value += (1 << k) * 2;
                }
            }

            return value + 2;
        };

        for(int i = 0; i < queries.size(); i++) {
            int u = queries[i][0];
            int w = queries[i][1];

            int val = lca(u, w);

            if(val == 0) {
                answer[i] = 0;
            }
            else {
                answer[i] = Pow(val-1);
            }
        }

        return answer;
    }
};