문자열 지옥에 빠진 호석

문제 설명

trie 문제집으로 분류되어 있기는 하지만, 구현 문제다. 그리고 굳이 trie를 안 써도 풀린다...

 

아무튼 해당 문제에서 특이한 점은 주어진 행렬이 이어져 있다고 생각해야 한다는 점이다.

예를 들어서 (0,0)에서 한 칸 위로 올라가면 (-1,0)이 된다. 그렇다면 해당 좌표는 탐색하지 않는 것이 옳지만, 행렬이 이어져 있다고 생각해야 하니 (2,0)을 탐색해줘야 한다.

즉, 좌표가 0보다 작으면 n-1 혹은 m-1의 위치로 들어가고 n 혹은 m의 위치가 되면 0으로 들어가야 한다는 것이다. 그리고 대각선도 탐색해야 한다는 것을 잊어서는 안 된다.

 

그리고 출력 부분도 조심해야 한다. 문제에서는 신이 좋아하는 문자열은 중복될 수 있다고 했다. 그러므로 중복된 문자열 또한 해당 문자열의 결과 값을 중복해서 출력해야 하는 것이다. 나는 이 부분을 놓쳐서 혼란과 당혹의 지옥에 빠졌었다. 분명 다 했는데 왜 안될까 싶을 때는 출력 조건과 제한을 살펴봅시다...

 

문제 풀이

첫 번째 접근

먼저 나는 trie를 이용하여서 문자열을 검색했다. 근데 이렇게 하지 말고 그냥 리스트에 넣어서 in해도 괜찮다.

그리고 새로운 좌표의 위치를 구해주는 방법은 좌표가 음수가 되면 n-1 혹은 m-1에서 맞는 값을 넣어주고 좌표가 n이나 m 이상일 때는 0으로 맞춰줬다. 이걸 더하고 나눠서 구하는 방법도 있는데, 나는 그냥 if else문 쓰는 게 더 직관적이라 좋아한다.

 

아, 참고로 행은 행대로 열은 열대로 조건을 비교해서 구해줘야 한다.

그리고 tmp 에다가 현재의 word를 넣어주고 새로운 좌표의 값을 word에 붙여서 새로운 문자열을 구해준다. 그리고 탐색하면 된다. 후에 탐색이 끝나면 word에다가 tmp를 넣어서 값을 이전의 것으로 맞춰준다. 혹은 굳이 함수 내에서 word를 변경하지 않고 재귀 함수로 넘어가는 인자에다가 맞춰줘도 된다. 이 부분은 후자의 것이 더 직관적이라고 생각한다.

 

그리고 출력을 해주는데, 여기서 중요한 건 중복되는 단어의 값은 중복해서 해당 정답을 출력해야 한다는 점이다. 

 

코드

import sys


class Trie:
    def __init__(self):
        self.root = {}

    def insert(self, word: str) -> None:
        cur = self.root
        for w in word:
            if w not in cur:
                cur[w] = {}
            cur = cur[w]
        cur['*'] = word

    def findWord(self, word: str) -> bool:
        cur = self.root
        for w in word:
            if w not in cur:
                return False
            cur = cur[w]
        return '*' in cur


N, M, K = map(int, sys.stdin.readline().rstrip().split(" "))
graph = [['' for _ in range(M)] for _ in range(N)]
for i in range(N):
    input1 = sys.stdin.readline().rstrip()
    for j in range(M):
        graph[i][j] = input1[j]
godWords = {}
godWordsList = []
t = Trie()
for i in range(K):
    godWord = sys.stdin.readline().rstrip()
    godWords[godWord] = 0
    godWordsList.append(godWord)
    if not t.findWord(godWord):
        t.insert(godWord)


dR = [-1, 0, 1, 0, -1, 1, -1, 1]
dC = [0, 1, 0, -1, -1, 1, 1, -1]
answer = 0


def dfs(word, r, c, cnt):
    if cnt > 5:
        return

    if word in godWords:
        godWords[word] += 1
        #return

    for i in range(8):
        nR = r + dR[i]
        nC = c + dC[i]

        if nR < 0:
            nR = N-1
        elif nR >= N:
            nR = 0

        if nC < 0:
            nC = M-1
        elif nC >= M:
            nC = 0

        tmp = word
        word += graph[nR][nC]
        dfs(word, nR, nC, cnt+1)
        word = tmp


for i in range(N):
    for j in range(M):
        dfs(graph[i][j], i, j, 1)

for w in godWordsList:
    if w in godWords:
        print(godWords[w])