프로그래머스 문제 - 가사 검색

프로그래머스 가사 검색 문제 풀이를 진행하였습니다.

문제 해석

문제를 읽으면 아래와 같은 해석이 가능합니다.

특정 키워드가 몇개 포함되어 있는지 확인해야 합니다.
"fro??"는 "frodo", "front", "frost" 등 단어가 같고 ?수 만큼 단어가 배치된 키워드에 매치되며, "frame", "frozen"에는 매치되지 않습니다.
가사에 사용된 모든 단어들이 담긴 배열 words와 찾고자 하는 키워드가 담긴 배열 queries가 주어질 때, 각 키워드 별로 매치된 단어가 몇 개인지 순서대로 배열에 담아 반환해야 합니다.
검색 키워드는 와일드카드 문자인 '?'가 하나 이상 포함돼 있으며, '?'는 각 검색 키워드의 접두사 아니면 접미사 중 하나로만 주어집니다.
"frodo"('?'가 없음), "fr?do"('?'가 중간에 있음), "?ro??"('?'가 양쪽에 있음)같은 키워드는 등장하지 않습니다.

문제 풀이

단어의 특정 지점에서의 뒤에 몇 개의 단어가 존재하는지 알아야 하며, 또한 뒤에서부터 시작해서 특정 지점의 앞에 몇 개의 단어가 존재하는지도 알 수 있어야 합니다.
이러한 문제를 해결하는데는 트라이(Trie)알고리즘이 좋습니다.

일단 각각의 단어들을 트리처럼 저장하기 위해 클래스를 만들어줍니다.
접두사, 접미사 두 곳에서의 경우의 수가 존재할 수 있기에 두 개의 클래스를 만들어줍니다.

class word
{
    public:
        int nextWordCount[26] = {0, };
        word* child[26] = {nullptr, };
        int count = 0;
};
class backword
{
    public:
        int prevWordCount[26] = {0, };
        backword* child[26] = {nullptr, };
        int count = 0;
};

이후 주어진 단어들을 트리 형태로 만들어줍니다.
각 검색 키워드의 단어 수의 따라 찾는 방식이 다르니 검색 이전에 단어의 수를 알 수 있도록 클래스를 배열로 선언해줍니다.

단어로 된 트리를 완성시킨 후 주어진 키워드에 부합하는 단어의 수를 차례대로 vector에 넣어주면 됩니다.

전체 코드

#include <bits/stdc++.h>
#include <string>
#include <vector>

using namespace std;

class word
{
    public:
        int nextWordCount[26] = {0, };
        word* child[26] = {nullptr, };
        int count = 0;
};
class backword
{
    public:
        int prevWordCount[26] = {0, };
        backword* child[26] = {nullptr, };
        int count = 0;
};

word wordRoot[10001];
backword backwordRoot[10001];

void insertWord(string str)
{
    //접두사부터 시작해서 순서대로 단어 저장
    word* curNode = &wordRoot[str.size()];
    for(int i = 0; i < str.size(); i++)
    {
        if(curNode->child[str[i] - 'a'] != nullptr)
        {
            curNode->nextWordCount[str[i] - 'a']++;
            curNode->count++;
            curNode = curNode->child[str[i] - 'a'];
        }else
        {
            curNode->nextWordCount[str[i] - 'a']++;
            curNode->count++;
            curNode->child[str[i] - 'a'] = new word;
            curNode = curNode->child[str[i] - 'a'];
        }
    }
    //접미사부터 시작해서 순서대로 단어 저장
    backword* curBackNode = &backwordRoot[str.size()];
    for(int i = str.size()-1; i >= 0; i--)
    {
        if(curBackNode->child[str[i] - 'a'] != nullptr)
        {  
            curBackNode->prevWordCount[str[i] - 'a']++;
            curBackNode->count++;
            curBackNode = curBackNode->child[str[i] - 'a'];
        }else
        {
            curBackNode->prevWordCount[str[i] - 'a']++;
            curBackNode->count++;
            curBackNode->child[str[i] - 'a'] = new backword;
            curBackNode = curBackNode->child[str[i] - 'a'];
        }
    }
}

int findWordCount(string str)
{
    if(str[0] != '?')
    {
        word* curNode = &wordRoot[str.size()];
        if(curNode->count == 0)
            return 0;
        
        //다음 위치의 단어가 ?가 될때까지 탐색 이후 갯수 return
        for(int i = 0; i < str.size()-1; i++)
        {
            if(str[i+1] != '?')
            {
                if(curNode->child[str[i] - 'a'] != nullptr)
                    curNode = curNode->child[str[i] - 'a'];
                else
                    return 0;
            }else
                return curNode->nextWordCount[str[i] - 'a'];
        }
    }else{
        backword* curBackNode = &backwordRoot[str.size()];
        if(curBackNode->count == 0)
            return 0;
        //접두사와 접미사가 ?라면 return 0;
        if(str[str.size()-1] == '?')
        {
            return curBackNode->count;
        }
        //이전 위치의 단어가 ?가 될때까지 탐색 이후 갯수 return
        for(int i = str.size()-1; i > 0; i--)
        {
            if(str[i-1] != '?')
            {
                if(curBackNode->child[str[i] - 'a'] != nullptr)
                    curBackNode = curBackNode->child[str[i] - 'a'];
                else
                    return 0;
            }else
                return curBackNode->prevWordCount[str[i] - 'a'];
        }
    }
}

vector<int> solution(vector<string> words, vector<string> queries) {
    vector<int> answer;
    
    for(string str : words)
    {
        insertWord(str);
    }
    for(string str : queries)
    {
        answer.push_back(findWordCount(str));
    }
    
    return answer;
}

출저

https://school.programmers.co.kr/learn/courses/30/lessons/60060