2020 카카오 공채: 가사 검색

[문제 링크]

코딩테스트 연습 - 가사 검색

---

트라이 자료구조를 사용하여 문자열을 빠르게 검색하는 문제인데 여러가지 고려해줘야할 것이 많아서 어려웠던 문제였다.

 

먼저 trie 객체에 count 변수를 추가하고 words 문자열이 각 노드를 방문하는 횟수를 카운팅함으로써 해당 위치에서 '?' 를 만났을 경우 가능한 가지수를 탐색하지 않고 바로 반환받을 수 있도록 하였다.

 

"????o" 와 같이 '?'가 접두사로 등장하는 경우에는 queries를 역순으로 탐색 해줘야하기 때문에 words를 뒤집은 문자열을 가지고 새로운 trie를 한개 더 만들었다.

 

마지막으로 '?'까지 문자열은 일치하지만 길이가 다른 경우를 걸러주기 위해 trie 객체를 words와 queries의 최대 길이인 10000만큼의 크기를 가진 배열로 선언하여 길이에 따라 다른 트라이에 담아주었다.

---

#include <iostream>
#include <cstring>
#include <string>
#include <algorithm>
#include <vector>
using namespace std;

const int ALPHABETS = 26;

int toNumber(char ch) 
{ 
	if (ch == '?') return ALPHABETS;
	
	return ch - 'a'; 
}

struct TrieNode
{
	TrieNode* children[ALPHABETS];
	int count;
	bool terminal;

	TrieNode() : terminal(false), count(1)
	{
		memset(children, 0, sizeof(children));
	}
	~TrieNode()
	{
		for (int i = 0; i < ALPHABETS; i++)
			if (children[i]) delete children[i];
	}

	void insert(string& key, int here)
	{
		if (here == key.size()) terminal = true;
		else
		{
			int next = toNumber(key[here]);

			//자식 노드가 없다면 생성
			if (children[next] == NULL)
				children[next] = new TrieNode();
			else
				children[next]->count++;

			children[next]->insert(key, here + 1);
		}
	}

	int find(string& key, int here)
	{
		if (here == key.size()) return this->terminal ? 1 : 0;

		int next = toNumber(key[here]);

		int cnt = 0;

		if (next == ALPHABETS)
		{
			for (int i = 0; i < ALPHABETS; i++)
				if (children[i] != NULL)
					cnt += children[i]->count;
		}
		else
		{
			if (children[next] == NULL) return 0;
			else
			{
				cnt = children[next]->find(key, here + 1);
			}
		}

		return cnt;
	}
};

vector<int> solution(vector<string> words, vector<string> queries) {
	vector<int> answer;

	vector<TrieNode> word(10001);
	vector<TrieNode> reword(10001);

	for (int i = 0; i < words.size(); i++)
	{
		int len = words[i].size();

		word[len].insert(words[i], 0);

		reverse(words[i].begin(), words[i].end());
		
		reword[len].insert(words[i], 0);
	}

	for (int i = 0; i < queries.size(); i++)
	{
		int len = queries[i].size();

		if(queries[i][0] != '?')
			answer.push_back(word[len].find(queries[i], 0));
		else
		{
			reverse(queries[i].begin(), queries[i].end());
			answer.push_back(reword[len].find(queries[i], 0));
		}
	}
	return answer;
}

알고리즘 200일 프로젝트 - 144 day