[프로그래머스] 가사 검색
#include <iostream>
#include <string>
#include <map>
#include <vector>
#include <algorithm>
using namespace std;
const int ALPHABET_SIZE = 26;
struct TrieNode {
struct TrieNode *children[ALPHABET_SIZE];
int numOfLeaf;
};
struct TrieNode *getNode(void) {
struct TrieNode *pNode = new TrieNode;
for (int i = 0; i < ALPHABET_SIZE; i++)
pNode->children[i] = NULL;
pNode->numOfLeaf = 0;
return pNode;
}
void insert(struct TrieNode *root, string key) {
struct TrieNode *pCrawl = root;
for (int i = 0; i < key.length(); i++){
pCrawl->numOfLeaf++;
int index = key[i] - 'a';
if (!pCrawl->children[index])
pCrawl->children[index] = getNode();
pCrawl = pCrawl->children[index];
}
}
int search(struct TrieNode *root, string key) {
struct TrieNode *pCrawl = root;
for (int i = 0; i < key.length(); i++){
if (key[i] == '?') {
return pCrawl->numOfLeaf;
}
int index = key[i] - 'a';
if (!pCrawl->children[index]) return 0;
pCrawl = pCrawl->children[index];
}
return 1;
}
vector<int> solution(vector<string> words, vector<string> queries) {
words.erase(unique(words.begin(), words.end()), words.end());
map<int, TrieNode*> trieRoot;
map<int, TrieNode*> reverseTrieRoot;
for (int i = 0; i < words.size(); i++) {
int len = words[i].length();
if (trieRoot.find(len) == trieRoot.end()) {
struct TrieNode *root = getNode();
trieRoot[len] = root;
struct TrieNode *reverseRoot = getNode();
reverseTrieRoot[len] = reverseRoot;
}
insert(trieRoot[len], words[i]);
reverse(words[i].begin(), words[i].end());
insert(reverseTrieRoot[len], words[i]);
}
vector<int> answer;
for (int i = 0; i < queries.size(); ++i) {
int len = queries[i].length();
if (trieRoot.find(len) == trieRoot.end()) {
answer.push_back(0);
continue;
}
if (queries[i][0] == '?') {
reverse(queries[i].begin(), queries[i].end());
answer.push_back(search(reverseTrieRoot[len], queries[i]));
}
else {
answer.push_back(search(trieRoot[len], queries[i]));
}
}
return answer;
}
int main() {
solution({ "frodo", "front", "frost", "frozen", "frame", "kakao" }, { "frod?", "????o", "fr???", "fro???", "pro?", "?????" });
}