개요
이 글에서는 C++에서 BIT(Binary Indexed Tree, 펜윅 트리)를 활용하여 색상이 칠해진 트리의 특정 서브트리에 포함된 서로 다른 색상의 개수를 효율적으로 구하는 프로그램을 다룹니다.
입력으로는 각 노드가 배열로 주어진 색상을 가지는 루트 있는 트리(rooted tree)가 제공됩니다. 우리의 목표는 주어진 노드를 루트로 하는 서브트리 안에 존재하는 서로 다른 색상 노드의 개수를 구하는 것입니다.
접근 방법
매 쿼리마다 서브트리를 직접 순회하면 최악의 경우 O(N×Q)의 시간이 소요되어 비효율적입니다. 이 문제는 다음과 같은 기법들을 조합하여 더 빠르게 해결할 수 있습니다.
- 오일러 투어(Euler Tour): DFS를 수행하면서 각 노드의 진입 시간(in-time)과 퇴장 시간(out-time)을 기록합니다. 어떤 노드 v의 서브트리는 평탄화된 배열에서 [in-time(v), out-time(v)] 구간에 해당합니다.
- 오프라인 쿼리 처리: 쿼리를 시작 시간 순으로 정렬한 뒤, 포인터를 한 번씩만 앞으로 이동시키며 처리합니다.
- 중복 색상 제거: 각 색상별로 현재 위치 이전의 마지막 등장 위치만 BIT에 1로 표시하고, 새로운 등장 위치가 나타나면 이전 위치를 제거합니다.
- BIT(펜윅 트리): 누적 합을 O(log N)에 계산하여 [in-time, out-time] 범위 내의 서로 다른 색상 개수를 빠르게 구합니다.
구현 코드
#include<bits/stdc++.h>
#define MAXIMUM_COLOUR 1000005
#define MAXIMUM_NUMBER 100005
using namespace std;
vector<int> tree[MAXIMUM_NUMBER];
vector<int> table[MAXIMUM_COLOUR];
int isTraversing[MAXIMUM_COLOUR];
int bit[MAXIMUM_NUMBER], getVisTime[MAXIMUM_NUMBER],
getEndTime[MAXIMUM_NUMBER];
int getFlatTree[2 * MAXIMUM_NUMBER];
bool vis[MAXIMUM_NUMBER];
int tim = 0;
vector< pair< pair<int, int>, int> > queries;
// 각 쿼리의 결과 저장
int ans[MAXIMUM_NUMBER];
// BIT 업데이트 함수
void update(int idx, int val) {
while ( idx < MAXIMUM_NUMBER ) {
bit[idx] += val;
idx += idx & -idx;
}
}
// BIT로 누적 합 조회
int queryingTree(int idx) {
int result = 0;
while ( idx > 0 ) {
result += bit[idx];
idx -= idx & -idx;
}
return result;
}
// DFS를 수행하며 트리를 평탄화
void preformingDFS(int v, int color[]) {
// 노드 방문 처리
vis[v] = 1;
getVisTime[v] = ++tim;
getFlatTree[tim] = color[v];
vector<int>::iterator it;
for (it=tree[v].begin(); it!=tree[v].end(); it++)
if (!vis[*it])
preformingDFS(*it, color);
getEndTime[v] = ++tim;
getFlatTree[tim] = color[v];
}
// 트리에 간선 추가
void addingNewEdge(int u, int v) {
tree[u].push_back(v);
tree[v].push_back(u);
}
// 각 색상의 첫 등장 위치를 BIT에 표시
void markingFirstFind(int n) {
for (int i = 1 ; i <= 2 * n ; i++) {
table[getFlatTree[i]].push_back(i);
if (table[getFlatTree[i]].size() == 1) {
update(i, 1);
isTraversing[getFlatTree[i]]++;
}
}
}
// 시작 지점 기준으로 정렬된 쿼리를 순서대로 처리
void calcQuery() {
int j = 1;
for (int i=0; i<queries.size(); i++) {
for ( ; j < queries[i].first.first ; j++ ) {
int elem = getFlatTree[j];
update( table[elem][isTraversing[elem] - 1], -1);
if ( isTraversing[elem] < table[elem].size() ){
update(table[elem][ isTraversing[elem] ], 1);
isTraversing[elem]++;
}
}
ans[queries[i].second] = queryingTree(queries[i].first.second);
}
}
// 서로 다른 색상 노드 개수 계산
void calcAllColours(int color[], int n, int qVer[], int qn) {
preformingDFS(1, color);
for (int i=0; i<qn; i++)
queries.push_back(make_pair(make_pair(getVisTime[qVer[i]] , getEndTime[qVer[i]]), i) );
sort(queries.begin(), queries.end());
markingFirstFind(n);
calcQuery();
for (int i=0; i<queries.size() ; i++) {
cout << "All distinct colours in the given tree: " << ans[i] << endl;
}
}
int main() {
int number = 6;
int color[] = {0, 2, 3, 3, 4, 1};
addingNewEdge(1, 2);
addingNewEdge(1, 3);
addingNewEdge(2, 4);
int queryVertices[] = {3, 2};
int qn = sizeof(queryVertices)/sizeof(queryVertices[0]);
calcAllColours(color, number, queryVertices, qn);
return 0;
}예제 설명
위 예제에서 트리 구조는 다음과 같습니다.
- 노드 1(색상 2)은 노드 2(색상 3) 및 노드 3(색상 3)과 연결
- 노드 2(색상 3)는 노드 4(색상 4)와 연결
노드 3을 루트로 하는 서브트리에는 노드 3 하나만 존재하므로 고유한 색상은 1개입니다. 반면 노드 2를 루트로 하는 서브트리에는 색상 3과 색상 4를 가진 두 노드가 있으므로 고유한 색상은 2개입니다.
실행 결과
All distinct colours in the given tree: 1 All distinct colours in the given tree: 2
시간 복잡도
DFS 수행과 초기화에 O(N log N), 쿼리 처리는 포인터 이동과 BIT 연산으로 인해 전체 O((N + Q) log N)의 시간 복잡도를 가집니다. 이는 매 쿼리마다 서브트리를 직접 순회하는 O(N × Q) 방식보다 훨씬 효율적입니다.