[Swift] 프로그래머스 - 카카오 가사 검색

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

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

 

정확도가 25점, 효율성이 75점이네요 Trie로 풀었습니다.

Trie 배우고 감으로 구현했는데 돌아가서 다행이네요.. ㅎㅎ

import Foundation

class TrieNode<Value: Hashable> {
    var value: Value?
    weak var parent: TrieNode?
    var children = [Value: TrieNode]()
    var childrenCount = 0
    
    init(value: Value? = nil, parent: TrieNode? = nil) {
        self.value = value
        self.parent = parent
    }
}

class Trie {
    typealias Node = TrieNode<Character>
    
    let root: Node = Node()
    
    func insert(value: String) {
        var current: Node = root
        
        for char in value {
            current.childrenCount += 1
            if current.children[char] != nil {
                current = current.children[char]!
            } else {
                current.children[char] = Node(value: char, parent: current)
                current = current.children[char]!
            }
        }
    }
    
    func count(prefix: String) -> Int {
        var current: Node = root
        
        for char in prefix {
            if char == "?" { return current.childrenCount }
            
            if current.children[char] != nil { // 자식이 있으면
                current = current.children[char]!
            } else { // 자식이 없으면
                return 0
            }
        }
        
        return 0
    }
}

// Trie 트리
func solution(_ words:[String], _ queries:[String]) -> [Int] {
    var tries = [Int: Trie]()
    var reversedTries = [Int: Trie]()
    var result = [Int]()
    
    // 1. 단어 집어넣기
    for word in words {
        if tries[word.count] != nil {
            tries[word.count]!.insert(value: word)
        } else {
            tries[word.count] = Trie()
            tries[word.count]!.insert(value: word)
        }
        
        let reversedWord = String(word.reversed())
        if reversedTries[word.count] != nil {
            reversedTries[word.count]!.insert(value: reversedWord)
        } else {
            reversedTries[word.count] = Trie()
            reversedTries[word.count]!.insert(value: reversedWord)
        }
    }
    
    // 2. 단어 수 세기
    for query in queries {
        if !query.starts(with: "?") {
            if let trie = tries[query.count] {
                result.append(trie.count(prefix: query))
            } else {
                result.append(0)
            }
        } else {
            let reversedQuery = String(query.reversed())
            if let trie = reversedTries[query.count] {
                result.append(trie.count(prefix: reversedQuery))
            } else {
                result.append(0)
            }
        }
    }
    
    return result
}

// print(solution(["frodo", "front", "frost", "frozen", "frame", "kakao"], ["fro??", "????o", "fr???", "fro???", "pro?"])) // expects [3, 2, 4, 1, 0]