[Swift] 2020 KAKAO BLIND RECRUITMENT 가사 검색



Problem

 

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

---

Solution

 

해당 문제는 Trie 자료구조를 사용해서 풀어야 하는 문제입니다.

 

(Trie 관련해서는 따로 정리해서 올리도록 하겠습니다!ㅜ)

 

1. Trie 자료 구조를 만들어준다.

 

정석으로 만든건 아니고 간단하게 단어를 Trie 자료구조에 맞게 삽입하는 insert 메서드와

 

글자의 자식이 몇 개 있는지 세어주는 getCount 메서드를 구현했습니다.

 

class Node {
    var value:String
    var count:Int = 0
    var children:[String:Node] = [:]
    init(value:String) {
        self.value = value
    }
    func append(_ value:String) {
        self.children[value] = Node(value: value)
    }
}

class Trie {
    let root:Node = Node(value:"")
    
    func insert(_ word:String) {
        var current = root
        let map = word.map{String($0)}
        map.forEach {
            current.count += 1
            if !current.children.keys.contains($0) {
                current.append($0)
            }
            current = current.children[$0]!
        }
    }
    
    func getCount(_ word:String) -> Int {
        var current = root
        let map = word.map{String($0)}
        for w in map {
            if !current.children.keys.contains(w) {
                return 0
            }
            current = current.children[w]!
        }
        return current.count == 0 ? 1 : current.count
    }
}

 

2. 올바른 Trie와 뒤집어진 글자를 담을 Trie를 만들어준다.

 

글자로 시작하는 단어는 올바른 trie로, ?로 시작하는 단어는 뒤집어진 trie인 reverse로 따로 관리해줍니다.

 

struct Tries {
    var trie:Trie,reverse:Trie
}

 

3. 단어 갯수에 맞게 분류해준다.

 

func classifyByWordLength(_ words:[String]) -> [Int:Tries] {
    var lengthDic:[Int:Tries] = [:]
    for word in words {
        if lengthDic[word.count] == nil {
            let tries = Tries(trie: Trie(), reverse: Trie())
            tries.trie.insert(word)
            tries.reverse.insert(String(word.reversed()))
            lengthDic[word.count] = tries
        }else {
            lengthDic[word.count]!.trie.insert(word)
            lengthDic[word.count]!.reverse.insert(String(word.reversed()))
        }
    }
    return lengthDic
}

 

4. 쿼리에 맞게 갯수를 담아준다.

 

주의할 점은 query 안의 단어의 길이가 lengthDic에 없다면 0을 반환하는 것입니다.

 

func queriesResults(_ queries:[String],_ lengthDic:[Int:Tries]) -> [Int] {
    var results:[Int] = []
    for query in queries {
        let count = query.count
        if lengthDic[count] == nil {
            results.append(0)
            continue
        }
        var pre = query.map{String($0)}.filter{$0 != "?"}.joined()
        pre = query.first == "?" ? String(pre.reversed()) : pre
        let trie = query.first == "?" ? lengthDic[count]!.reverse : lengthDic[count]!.trie
        let result = trie.getCount(pre)
        results.append(result)
    }
    return results
}

 

5. results를 반환한다.

 

func solution(_ words:[String], _ queries:[String]) -> [Int] {
    let lengthDic:[Int:Tries] = classifyByWordLength(words)
    return queriesResults(queries,lengthDic)
}

---

Source Code