코딩 테스트 - 트리 트리오 중간값

진행한 코딩 테스트 문제를 기록하기 위한 글 입니다.

해당 문제 링크 : https://school.programmers.co.kr/learn/courses/30/lessons/68937

문제 설명

n개의 노드와 n-1개의 간선을 갖는 트리 구조에서, 어떤 세 노드(a,b,c)에 대해 a에서 b까지의 거리, b에서 c까지의 거리, c에서 a까지의 거리를 각각 구하고, 그 세개의 거리중 중간값을 찾아내는 문제이다.

이때 이 중간값은 전달받은 트리 구조에서 가능한 가장 큰 값이어야 한다.

 
이 문제를 보고 처음에 생각해낸건 가장 길게 연결된 간선과 그 끝의 두 노드를 찾게되면 그 두 노드는 정답인 a,b,c에 무조건적으로 포함이 되나? 라는 의문이었다.

간선의 수가 n-1개라는 조건하에 여러 형태의 그래프를 생각하고 실제로 배치도 해보면서 테스트를 진행해본 결과, 가장 긴 간선 조합 및 그에 의한 두 노드와 나머지 노드 하나를 사용하는게 최대의 중간값을 얻을 수 있다는 결론에 도달했다.

따라서, 가장 먼저 가장 긴 경로를 찾는 방법을 생각해보게 되었는데, 처음엔 뭔가 기준이 있어야만 가능한게 아닌가? 하는 생각만 들었다.

문제를 다시 보면 문제의 이미지는 트리보단 그래프였는데, 문제를 다시 읽어보면 이는 분명한 트리로, 지금까지 머릿속에서 인지하던 트리와는 다르지만 분명한 트리 구조라는 것을 인지하게 되었다.

이는 즉, 이어짐 없이 따로 동떨어진 노드도 없고 사이클 구조도 존재하지 않는다는 것을 의미하며, 사이클이 존재하지 않는다면 간선이 하나만 연결된 노드가 존재하고, 결국 최장 경로는 이 하나만 연결된 노드끼리 연결되어야 한다는 부분까지 접근하게 되었다.

뿌리냐 잎사귀냐 같은 생각을 진행하다가, 문제를 보고 다시 한번 기존에 인지하던 단순한 트리와는 달리 양방향으로 접근이 가능한 트리라는 것을 깨닫게 되었으며, 이런 양방향 트리 구조에서의 최장 경로를 찾는 정형화 된 무언가가 무조건 존재할거라는 생각이 들었다.

이 생각을 토대로 트리의 지름이라는 개념을 찾게 되었는데, 지름은 곧 계속 생각했던 대로 가장 긴 간선의 경로 끝 두 노드와 그 경로를 의미했고, 구하는 방법으로는 아무런 노드에서부터 시작해서 가장 멀리떨어진 노드를 하나 찾고, 거기서부터 한번 더 가장 멀리떨어진 노드를 찾는 간단한 방법으로 구할 수 있었다.

구현 코드

코드는 크게 그래프 구현 및 BFS 탐색 메소드로 나뉘어진다.

using System;
using System.Collections.Generic;

public class Solution
{
    bool[] visited; //방문 배열
    List<int>[] graph; //그래프 리스트
    Queue<(int node, int distance)> queue = new Queue<(int node, int distance)>(); //노드와 해당 노드까지의 거리를 저장해두는 큐

    public int solution(int n, int[,] edges)
    {
        //그래프 리스트 배열 크기 설정, 노드 숫자를 있는 그대로 인덱스로 사용
        graph = new List<int>[n + 1];
        visited = new bool[n + 1];

        //배열 초기화
        for (int i = 0; i <= n; i++)
        {
            graph[i] = new List<int>();
        }

        //그래프 연결
        for (int i = 0; i < edges.GetLength(0); i++)
        {
            int from = edges[i, 0];
            int to = edges[i, 1];

            //양방향으로 연결 진행
            graph[from].Add(to);
            graph[to].Add(from);
        }

        //첫번째 a 노드 확보
        int a = BFS(1, 0, out int aaa);

        //두번째 b 노드 확보
        int b = BFS(a, 0, out int aToB);

        //a -> c 예정, b -> c 예정 노드 확보
        int c_1 = BFS(a, b, out int aToc);
        int c_2 = BFS(b, a, out int bToc);

        //두 노드 중 거리값이 더 큰 노드를 c로 저장
        int c;
        if(aToc >= bToc)
        {
            //만약 a -> c가 맞다면
            c = c_1;

            //b 노드에서 c 노드까지의 거리 확보
            bToc = GetTargetDistanceBFS(b, c);
        } 
        else
        {
            //반대로 b -> c가 맞다면
            c = c_2;

            //a 노드에서 c 노드까지의 거리 확보
            aToc = GetTargetDistanceBFS(a, c);
        }

        //최장거리를 제외한 두 거리중에서 더 큰 값이 중간값
        return Math.Max(aToc, bToc);
    }
}

메인 코드부터 살펴보자면, 그래프 구현 및 BFS 로직을 위한 각종 변수들부터 선언을 진행하고, 입력받은 n 값에 맞춰서 초기화를 진행해주도록 되어있다.

이후, 그래프를 양방향으로 구현해준 뒤, 지름을 찾기위한 BFS 로직이 실행된다.

int BFS(int startNode, int dontNode, out int returnDistance)
{
    //방문 배열, 큐 초기화
    Array.Clear(visited, 0, visited.Length);
    queue.Clear();

    //최종적으로 반환할 노드와 그 노드까지의 거리
    int returnNode = 0;
    returnDistance = 0;

    //시작 노드를 큐에 넣고, 방문 여부 저장
    queue.Enqueue((startNode, 0));
    visited[startNode] = true;

    while (queue.Count > 0)
    {
        //큐에 먼저 들어갔던 노드를 가져옴
        var current = queue.Dequeue();

        //노드와 그 노드까지의 거리 확보
        int node = current.node;
        int distance = current.distance;

        if(dontNode != node)
        {
            returnNode = node;
            returnDistance = distance;
        }

        //현재 노드와 관련된 노드들 중 아직 방문하지 않은 노드를 큐에 등록
        foreach (int next in graph[node])
        {
            if (!visited[next])
            {
                //방문 여부를 저장하면서 거리를 늘려 큐에 등록 진행
                visited[next] = true;
                queue.Enqueue((next, distance+1));
            }
        }
    }

    //마지막으로 체크한 노드 확보
    return returnNode;
}

매개변수는 총 3개로 시작 노드와 무시할 노드, 그리고 반환할 최종 거리를 의미한다.

BFS 탐색에 사용되는 큐는 노드의 번호와 시작 노드에서부터 그 노드까지의 거리의 조합으로 이루어져 있기 때문에, 탐색 전 시작 노드를 넣을 땐 그 거리를 0으로 잡아주면서 큐에 넣어주며 탐색을 시작한다.

큐에서 노드를 하나 꺼내고, 다음 탐색을 위해 현재 노드와 연결된 다른 노드를 찾으면서, 새롭게 찾게된 노드를 큐에 다시 넣어주는 방식으로 탐색을 진행하는데, 이때 중요한건 한단계 더 떨어진 노드를 찾은 것이기에 거리에 +1을 진행해주는 것이다.

이 과정에서 새롭게 찾은 노드를 계속해서 최종적으로 반환할 노드로 업데이트 해주는데, 만약 현재 노드를 무시해야 하는 상황에서는 반환용 노드로 업데이트 하지 않는다.

이런 BFS 탐색 메소드를 이용해서 1번 노드부터 가장 멀리 떨어진 노드를 찾아 a 노드로 지정하고, a 노드에서 가장 멀리 떨어진 노드를 찾아 b 노드로 지정함과 동시에 a->b 까지의 거리를 확보한다.

즉, 가장 긴 간선의 조합과 그 노드를 찾게 된 것으로 이제 나머지 c 노드를 찾아주어야 한다.

 
c 노드는 특이하게 a와 b 두 노드에서 각각 BFS 탐색을 진행한다.

이때 아무런 제한 없이 탐색을 진행하면 결국 a에서는 항상 b 노드를, b에서는 항상 a노드를 반환할 가능성이 존재하기에 무시할 노드 매개변수를 이용하여 a에서 시작할때는 b를 무시하고, b에서 시작할때는 a를 무시해준다.

이후 그렇게 구한 두 노드와 그 간선 거리중에 더 큰쪽을 실제 c 노드로 지정하고, 확정이 되면 남은 경로에 대해 추가적인 BFS 탐색을 진행한다.

이렇게 c 예정 노드를 a와 b 양쪽에서 각각 구해서 판별하지 않으면 자칫 엉뚱한 노드가 c 노드가 될 확률이 존재하기 때문이다.

예를 들어서 1 - 2 - 3 - 4 - 5 구조에서 4 - 6까지 이어진 구조를 생각해보자.

맨 처음 a 노드를 구할 때 1번 노드부터 시작하기 때문에 a노드는 5 아니면 6이 되는데, 일단은 5라고 생각하고 거기서부터 가장 멀리 떨어진 b 노드도 구하면 b 노드는 1이 된다.

이때 b 노드에서부터 c 노드를 구하면 5를 제외하고 가장 멀리 떨어진 6이 c 노드가 될 테지만, a에서 c 노드를 구하면 1을 제외하고 가장 멀리 떨어진 2가 c 노드가 되어버린다.

이 경우 b노드에서 시작한 c 노드가 정답인데, 만약 6이 4가 아니라 2에 붙어있었다면 반대로 a노드에서 시작한 c 노드가 6이 되면서 그쪽이 정답이 되어버린다.

따라서 두 노드에서 각각 서로를 제외하고 가장 멀리 떨어진 노드를 구한 다음, 떨어진 거리를 체크해서 더 큰 쪽을 실제 c 노드로 삼고, c와 나머지 노드의 거리를 구해준다.

//시작과 타겟 노드가 존재하는 BFS 로직
int GetTargetDistanceBFS(int startNode, int targetNode)
{
    //방문 배열, 큐 초기화
    Array.Clear(visited, 0, visited.Length);
    queue.Clear();

    queue.Enqueue((startNode, 0));
    visited[startNode] = true;

    while (queue.Count > 0)
    {
        //큐에 먼저 들어갔던 노드를 가져옴
        var current = queue.Dequeue();

        //노드와 그 노드까지의 거리 확보
        int node = current.node;
        int distance = current.distance;

        //만약 현재 노드가 찾는 노드라면
        if (targetNode == node)
        {
            return distance;
        }

        //현재 노드와 관련된 노드들 중 아직 방문하지 않은 노드를 큐에 등록
        foreach (int next in graph[node])
        {
            if (!visited[next])
            {
                //방문 여부를 저장하면서 거리를 늘려 큐에 등록 진행
                visited[next] = true;
                queue.Enqueue((next, distance + 1));
            }
        }
    }

    //타겟 노드를 못찾으면 그냥 0 반환
    return 0;
}

이번엔 시작 노드와 목적 노드가 명확하기 때문에, 전체적인 탐색 과정은 동일하지만 그 과정에서 타겟 노드를 찾으면 타겟노드까지의 거리를 반환해주는 메소드를 구현해주었다.

이 과정들을 통해 a,b,c 세 노드와 각 노드간의 거리를 전부 파악할 수 있게 되었다.

나머지는 세 거리중 중간값을 반환해주는 부분만 남았는데, 우선 a에서부터 b까지의 경로는 확정적으로 가장 긴 경로이기 때문에 [a,c], [b,c] 두 경로만 체크해주면 된다.

가장 큰 값이 빠진 나머지 두 값중에서 더 큰 값이 바로 세 값중에 중간값이 되기 때문에 이는 Max 기능을 이용해서 남은 두 경로중 더 큰 값을 반환하며 로직이 종료된다.

 

사실 남은 경로의 길이를 구하는 과정과 그 메소드인 GetTargetDistanceBFS 메소드는 문제의 정답을 제출할 뿐이라면 필요가 없다.

int c_1 = BFS(a, b, out int aToc);
int c_2 = BFS(b, a, out int bToc);

c 예정 노드를 구하던 부분 중에서 aToc가 5, bToc가 4라고 가정해보자.

거리가 더 먼 만큼 실제로 c_1 노드가 정답이 되는데, 이때 남은 c_2 노드까지의 거리인 bToc의 4가 의미하는 바는 b에서 시작해서 a를 제외한 모든 노드들 중에서 가장 먼 거리는 4라는 의미이다.

이때 같이 찾은 c_1 노드 또한 a가 아닌 노드 중 하나이기에 결국 b에서 c_1까지의 거리도 4 이하가 될 수 밖에 없다.

따라서, b에서 c 노드의 거리를 새로 구할 필요 없이 가장 긴 지름 거리를 제외하고 남은 두 거리중에서 더 큰 aToc의 값이 중간값이 된다.

물론 b에서부터 a를 제외하고 찾은 c_2는 c_1 노드와는 다른 노드일 가능성이 존재하지만, 결국 문제에서 원하는건 실제 c 노드가 아니라 거리이기에 굳이 c 노드를 하나로 확정하고 새롭게 탐색을 진행해주지 않아도 문제풀이 자체는 가능하며, 그에 따라서 실제 로직도 더욱 생략할 수는 있었다.

그렇지만 실제로 머릿속에서 생각한 흐름에 맞추기 위해 c 노드를 하나로 확정하고 c 노드까지의 거리를 확정하며 비교하는 방식으로 문제를 해결해주었다.