C++ / 백준 / 14725 / 개미굴

문제

 

https://www.acmicpc.net/problem/14725

14725번: 개미굴

문제풀이

 

트라이 문제이다.

입력으로는 알파벳 대문자로만 이루어진 string이 들어온다.

 

트라이는 map을 이용하여 자녀 트라이를, 깊이를 나타내는 depth, 문자열의 끝 여부를 나타내는 isEnd로 구성하였다.

 

	void insert(vector<string>& v) {
		Trie* now = this;
		for (int idx = 0; idx < v.size(); ++idx) {
			if (now->child.find(v[idx]) == now->child.end()) {
				now->child[v[idx]] = new Trie();
			}
			now->child[v[idx]]->depth = now->depth + 1;
			now = now->child[v[idx]];

			if (idx + 1 == v.size()) now->isEnd = true;
		}
	}

문자열을 입력하는 함수이다.

루트 트라이부터 시작하여 문자열 벡터를 탐색한다.

자녀 트라이에 해당 문자열이 없다면 새로 자녀 트라이를 생성한다.

자녀 트라이의 깊이는 부모 트라이의 깊이 +1로 설정한다.

그다음 현재 위치를 자녀 트라이로 이동한다.

현재 문자열이 끝이라면 isEnd를 true로 바꾼다.

 

 

	void print() {
		Trie* now = this;
		for (auto iter = now->child.begin(); iter != now->child.end(); ++iter) {
			for (int idx = 0; idx < now->depth; ++idx) cout << "--";
			cout << iter->first << '\n';
			iter->second->print();
		}
	}

루트 트라이부터 시작하여 모든 자녀의 트라이를 탐색한다.

해당 트라이의 깊이만큼 "--"를 출력한 뒤, 현재 문자열을 출력한다.

그다음 자녀 트라이로 이동하여 반복한다.

 

코드

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

struct Trie {
	map<string, Trie*> child;
	int depth;
	bool isEnd;

	Trie() : child(), depth(0), isEnd(false) {}

	void insert(vector<string>& v) {
		Trie* now = this;
		for (int idx = 0; idx < v.size(); ++idx) {
			if (now->child.find(v[idx]) == now->child.end()) {
				now->child[v[idx]] = new Trie();
			}
			now->child[v[idx]]->depth = now->depth + 1;
			now = now->child[v[idx]];

			if (idx + 1 == v.size()) now->isEnd = true;
		}
	}
	void print() {
		Trie* now = this;
		for (auto iter = now->child.begin(); iter != now->child.end(); ++iter) {
			for (int idx = 0; idx < now->depth; ++idx) cout << "--";
			cout << iter->first << '\n';
			iter->second->print();
		}
	}
};

int main() {
	ios::sync_with_stdio(false);
	cin.tie(NULL);

	int N(0), K(0);
	cin >> N;
	Trie* trie = new Trie();
	for (int idx = 0; idx < N; ++idx) {
		cin >> K;
		vector<string>words(K, "");
		for (int idx2 = 0; idx2 < K; ++idx2) {
			cin >> words[idx2];
		}
		trie->insert(words);
	}
	trie->print();
}