C++ BFS 알고리즘으로 트리의 특정 레벨에 있는 노드 개수 구하기

트리의 노드들을 정점으로 가지는 무방향 그래프가 주어졌을 때, BFS(너비 우선 탐색) 알고리즘을 이용해 트리의 특정 레벨(level)에 있는 노드의 개수를 구하는 것이 목표입니다.

BFS 알고리즘이란?

BFS는 그래프나 트리를 레벨 단위로 순차적으로 탐색하는 알고리즘입니다. 레벨 0의 시작 노드에서 출발하여, 먼저 해당 노드와 직접 연결된 레벨 1의 모든 노드를 방문한 뒤, 이어서 다음 레벨의 노드들을 차례로 탐색합니다.

현재 레벨의 노드들을 수평 방향으로 탐색합니다.
이어서 다음 레벨의 노드들도 같은 방식으로 탐색합니다.

예제로 이해하기

예제 1

입력: level = 2

출력: BFS로 구한 해당 레벨의 노드 개수: 1

설명: 위 그래프에서 각 레벨에는 노드가 하나씩만 존재하므로, 레벨 2의 노드 개수는 1입니다.

예제 2

입력: level = 1

출력: BFS로 구한 해당 레벨의 노드 개수: 2

설명: 레벨 1에는 노드 1과 노드 2가 있으므로 개수는 2입니다.

해결 접근 방식

이 접근법에서는 각 노드를 순회하면서 자신의 부모 노드보다 한 레벨 더 깊은 값을 레벨로 설정합니다. 그래프는 인접 리스트(adjacency list) 방식으로 표현합니다.

예를 들어 시작 노드가 0이라면 다음과 같이 레벨이 결정됩니다.

level[0] = 0
level[1] = level[0] + 1 = 1, level[2] = level[0] + 1 = 1
level[3] = level[2] + 1 = 2, level[4] = level[2] + 1 = 2

알고리즘 단계

정점의 개수를 저장하는 data와 인접 리스트 포인터 next를 멤버로 가지는 node 클래스를 생성합니다.
공개 메서드 insert(int val, int point)는 그래프에 간선을 추가합니다. val을 point의 인접 리스트에 추가하고, point도 val의 인접 리스트에 추가합니다.
함수 count_nodes(int a, int b)는 시작 노드 a로부터 레벨 b에 있는 노드의 개수를 반환합니다.
초기 count 값을 0으로 설정합니다.
방문 여부를 저장할 bool 배열 check = new bool[data]를 선언합니다.
배열 arr[data]는 그래프의 각 정점 레벨을 저장합니다.
for 반복문으로 모든 정점을 미방문 상태로 초기화합니다(check[i] = false, arr[i] = 0).
BFS 탐색을 위한 큐 l1을 생성합니다.
시작 정점을 방문 처리(check[a] = true)하고, l1.push_back(a)로 큐에 추가한 뒤 레벨을 0으로 설정합니다(arr[a] = 0).
큐 l1이 빌 때까지 반복합니다.
front 요소를 꺼내고(l1.front(), l1.pop_front()) 해당 정점과 인접한 미방문 정점을 모두 방문 처리하여 큐에 추가합니다.
각 인접 정점의 레벨을 현재 정점 a의 레벨 + 1로 설정합니다.
while 반복문이 끝나면 for 반복문으로 arr[]를 순회하며 arr[i] == b인 경우 count를 증가시킵니다.
최종적으로 count를 결과로 반환합니다.

C++ 구현 코드

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

class node {
    int data;
    list < int > * next;
    public:
        node(int data) {
            this -> data = data;
            next = new list < int > [data];
        }
    void insert(int val, int point) {
        next[val].push_back(point);
        next[point].push_back(val);
    }
    int count_nodes(int a, int b);
};

int node::count_nodes(int a, int b) {
    int count = 0;
    bool * check = new bool[data];
    int arr[data];

    for (int i = 0; i < data; i++) {
        check[i] = false;
        arr[i] = 0;
    }

    list < int > l1;
    check[a] = true;
    l1.push_back(a);
    arr[a] = 0;

    while (!l1.empty()) {
        a = l1.front();
        l1.pop_front();
        for (auto it = next[a].begin(); it != next[a].end(); ++it) {
            if (!check[ * it]) {
                arr[ * it] = arr[a] + 1;
                check[ * it] = true;
                l1.push_back( * it);
            }
        }
    }
    for (int i = 0; i < data; i++) {
        if (arr[i] == b) {
            count++;
        }
    }
    return count;
}
int main() {
    node n1(5);
    n1.insert(1, 2);
    n1.insert(0, 3);
    n1.insert(1, 3);
    n1.insert(2, 4);

    int level = 1;

    cout << "BFS로 구한 해당 레벨의 노드 개수: " << n1.count_nodes(0, level);
    return 0;
}

실행 결과

BFS로 구한 해당 레벨의 노드 개수: 1

위 코드는 5개의 정점으로 구성된 트리를 만들고, 시작 노드 0에서 레벨 1에 있는 노드의 개수를 BFS로 계산하여 출력합니다. 시간 복잡도는 O(V + E)로, 정점과 간선의 수에 비례하여 효율적으로 동작합니다.