[1일 1알고] 가사 검색

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

프로그래머스

이번 문제는 Trie를 활용하는 문제입니다.

 

문제에서는 특정 단어들에 대해 검색을 할것이며 이 때 ?(와일드카드 문자)가 들어갈 수 있습니다.

다만 ?는 접두 혹은 접미사로만 붙을 수 있기 때문에 문제의 해결이 쉬워집니다.

 

일단 조건을 보았을 때 각 단어에 대해 무작정 찾는 것은 불가능할 것임을 알 수 있습니다.

 

이를 위해 Trie를 사용합니다.

 

Trie에 필요한 정보는 다음과 같습니다.

1. 현재 Trie Node에 들어가는 문자 ch

2. 현재 Trie Node 밑에 존재하는 경우의 수 count

3. 현재 Trie Node 밑에 존재하는 노드들로 이동하기 위한 포인터들

 

이렇게 구성한다면 만약에 far??라는 query가 들어왔을 때 far까지 탐색한 후 r에 존재하는 count를 가져오면 될 것입니다.

---

다만 Trie를 만들 때 고려해야하는 것들이 있습니다.

 

far??와 far???를 구별하기 위해서는 각 문자의 길이에 따른 개별 Trie를 만들어야합니다.

 

그리고 ??far와 같이 접두사로 ?가 왔을 때를 대비하기 위해 문자를 거꾸로 뒤집은 Trie를 만들어야합니다.

예를 들어 sofar라는 word가 있다면 rafos로 저장해야 ??far에 대한 query를 판별할 수 있을 것입니다.

---

 

using namespace std;

struct Node
{
    char ch = 0;
    long long cnt = 0;
    vector<Node*> nexts;
};

// idx = 현재 찾는 글자의 인덱스
void MakeTrie(int idx, string& word, Node* pCur)
{
    if (idx == word.size())
        return;
    int nextIdx = -1;
    for (int i = 0; i < pCur->nexts.size(); ++i)
    {
        if (pCur->nexts[i]->ch == word[idx])
        {
            nextIdx = i; // 다음 trie 노드
        }
    }

    Node* nextNode;

    if (nextIdx == -1)
    {
        nextNode = new Node;
        nextNode->ch = word[idx];
        nextNode->cnt = 1; 
        pCur->nexts.push_back(nextNode);
    }
    else
    {
        ++pCur->nexts[nextIdx]->cnt;
        nextNode = pCur->nexts[nextIdx];
    }
    MakeTrie(idx + 1, word, nextNode);
}


long long FindQueryCnt(int idx, string& query, Node* pCur, long long cnt)
{
    if (query[idx] == '?')
        return cnt;

    Node* nextNode = nullptr;
    for (int i = 0; i < pCur->nexts.size(); ++i)
    {
        if (pCur->nexts[i]->ch == query[idx])
        {
            cnt = pCur->nexts[i]->cnt;
            nextNode = pCur->nexts[i];
            break;
        }
    }
    if (nextNode == nullptr)
        return 0;
    else
    {
        return FindQueryCnt(idx + 1, query, nextNode, cnt);
    }
}
vector<int> solution(vector<string> words, vector<string> queries) {
    vector<int> answer;

    const int MAX_LEN = 10000;

    vector<Node*> tries(MAX_LEN + 1, nullptr);
    vector<Node*> reverseTries(MAX_LEN + 1, nullptr);

    for (string word : words)
    {
        int len = word.size();

        if (tries[len] == nullptr)
            tries[len] = new Node;

        ++tries[len]->cnt;
        MakeTrie(0, word, tries[len]);

        reverse(word.begin(), word.end());

        if (reverseTries[len] == nullptr)
            reverseTries[len] = new Node;

        ++reverseTries[len]->cnt;
        MakeTrie(0, word, reverseTries[len]);
    }

    for (string query : queries)
    {
        int len = query.size();

        if (query[0] == '?')
        {
            if (reverseTries[len] == nullptr)
            {
                answer.push_back(0);
                continue;
            }

            if (query[len - 1] == '?')
            {
                answer.push_back(reverseTries[len]->cnt);
            }
            else
            {
                reverse(query.begin(), query.end());
                long long cnt = FindQueryCnt(0, query, reverseTries[len], reverseTries[len]->cnt);
                answer.push_back(cnt);
            }
        }
        else
        {
            if (tries[len] == nullptr)
            {
                answer.push_back(0);
                continue;
            }

            long long cnt = FindQueryCnt(0, query, tries[len], tries[len]->cnt);
            answer.push_back(cnt);
        }
    }

    return answer;
}

 

정리하면 이와 같습니다.

 

재귀를 사용하는 함수가 들어갔기 때문에 디버깅을 해보면서 정확하게 구현했는지 체크하는게 중요해보였습니다.