Union-Find 연산에서 find함수 return할 때 주의할 점

Union-Find 연산에서 find()함수의 코드는 일반적으로 이렇게 짠다.

public static int find(int x){
    if(x==parent[x]) {
        return x;
    }
    return parent[x]=find(parent[x]);
}

 

❗️이때 return을 find(parent[x])이 아니고 왜 parent[x]=find(parent[x]);를 해야할까?

---

다음과 같은 트리 구조가 있다고 하자.

1 <- 2 <- 3 <- 4 <- 5

 

이런 경우 1이 루트 노드가 되고, parent 배열은 다음과 같다.

노드 Idx	1	2	3	4	5
parent[idx]	1	1	2	3	4

 

---

return find(parent[x]);일 경우

find(5)를 호출하면:

• find(5) -> find(4) -> find(3) -> find(2) -> find(1)
• 5번 재귀 호출 후 1을 반환
• parent 배열이 변하지 않는다.

---

return parent[x] = find(parent[x]);일 경우

find(5)를 호출하면:

• find(5) -> find(4) -> find(3) -> find(2) -> find(1)
• 5번 재귀 호출 후 1을 반환하지만, 돌아오면서 모든 노드의 부모를 1로 설정한다.
• parent 배열이 다음과 같이 변한다.

노드 Idx	1	2	3	4	5
parent[idx]	1	1	1	1	1

---

 

중요한 점은 parent 배열이 수정되어서 이 이후에 find(5)를 호출하면 바로 1을 반환하게 된다는 것이다.

 

return할 때 parent배열을 수정해주면 첫 번째 find 연산은 동일한 시간이 걸리지만, 이후의 모든 find 연산은 매우 빨라진다. 거의 상수 시간에 루트를 찾을 수 있게 된다. 이런 최적화는 특히 유니온-파인드 연산이 많이 반복되는 경우에 더 중요하다.

 

예를 들어 백준 18116번: 로봇 조합 문제는 find함수에서 위와 같이 최적화를 해주지 않으면 시간초과가 난다. 

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

public class _18116 {
    public static int[] parent;
    public static int[] num;
    public static void main(String[] args) throws IOException{
        BufferedReader br=new BufferedReader(new InputStreamReader(System.in));
        int N=Integer.parseInt(br.readLine());
        parent=new int[1000001];
        num=new int[1000001];
        Arrays.fill(num, 1);
        for(int i=0;i<=1000000;i++){
            parent[i]=i;
        }
        for(int i=0;i<N;i++){
            StringTokenizer st=new StringTokenizer(br.readLine());
            String instr=st.nextToken();
            if(instr.equals("I")){ //union
                int a=Integer.parseInt(st.nextToken());
                int b=Integer.parseInt(st.nextToken());
                union(a,b);
            }else{ //find
                int c=Integer.parseInt(st.nextToken());
                System.out.println(num[find(c)]);
            }
        }
    }
    public static void union(int x, int y){
        x=find(x);
        y=find(y);
        if(x!=y) {
            parent[y] = x;
            num[x]+=num[y];
        }
    }
    public static int find(int x){
        if(x==parent[x]) {
            return x;
        }
        return parent[x]=find(parent[x]);
    }

}