Created
February 14, 2018 18:23
-
-
Save gstraymond/56b041993c1de6b2490b54fdf0ebe86b to your computer and use it in GitHub Desktop.
Kotlin Trie
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package fr.gstraymond.search | |
data class Trie(private val char: Char, | |
private val indices: MutableList<Int> = mutableListOf(), | |
private val children: MutableList<Trie> = mutableListOf()) { | |
fun add(word: String, | |
index: Int) { | |
val first = word.first() | |
val trie = findOrCreate(first) | |
when (word.length) { | |
1 -> trie.indices.add(index) | |
else -> trie.add(word.drop(1), index) | |
} | |
} | |
fun get(word: String): List<Int> { | |
val first = word.first() | |
return find(first)?.run { | |
when (word.length) { | |
1 -> indices | |
else -> get(word.drop(1)) | |
} | |
} ?: listOf() | |
} | |
private fun findOrCreate(char: Char) = | |
find(char) ?: Trie(char).apply { this@Trie.children.add(this) } | |
private fun find(char: Char) = | |
children.find { it.char == char } | |
} | |
object TrieBuilder { | |
fun empty() = Trie(char = ' ') | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package fr.gstraymond.search | |
import io.kotlintest.matchers.shouldBe | |
import io.kotlintest.specs.StringSpec | |
class TrieSpec : StringSpec() { | |
init { | |
"empty trie should return nothing" { | |
TrieBuilder.empty().get("a") shouldBe listOf<Int>() | |
} | |
"hello should return 0" { | |
val emptyTrie = TrieBuilder.empty() | |
emptyTrie.add("hello", 0) | |
emptyTrie.get("hello") shouldBe listOf(0) | |
} | |
"multiple hello should return 0, 2" { | |
val emptyTrie = TrieBuilder.empty() | |
emptyTrie.add("hello", 0) | |
emptyTrie.add("world", 1) | |
emptyTrie.add("hello", 2) | |
emptyTrie.get("hello") shouldBe listOf(0, 2) | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment