72 lines
2.1 KiB
Kotlin
72 lines
2.1 KiB
Kotlin
package moe.kageru.spektacles
|
|
|
|
import kotlin.collections.LinkedHashMap
|
|
|
|
|
|
// maybe try this instead?
|
|
// https://arrow-kt.io/docs/recursion/intro/
|
|
sealed class Node<K, T : Any>(open val value: Any) {
|
|
class Leaf<K, T : Any>(override val value: List<T>) : Node<K, T>(value)
|
|
class Split<K, T : Any>(val children: Map<K, Node<K, T>>) : Node<K, T>(children)
|
|
class Value<K, T : Any>(override val value: T) : Node<K, T>(value)
|
|
|
|
fun <R : Any> map(op: (T) -> R): Node<K, R> {
|
|
return when (this) {
|
|
is Leaf -> Leaf(value.map(op))
|
|
is Value -> Value(op(value))
|
|
is Split -> Split(children.mapValues { it.value.map(op) })
|
|
}
|
|
}
|
|
|
|
fun groupBy(keyGenerator: (T) -> K): Node<K, T> {
|
|
return when (this) {
|
|
is Leaf -> Split(value.groupBy(keyGenerator).mapValues { Leaf<K, T>(it.value) })
|
|
is Value -> this
|
|
is Split -> Split(children.mapValues { it.value.groupBy(keyGenerator) })
|
|
}
|
|
}
|
|
|
|
override fun toString() = asString()
|
|
|
|
private fun asString(depth: Int = 0): String {
|
|
return when (this) {
|
|
is Leaf -> value.joinToString(indent(depth), prefix = "\n${indent(depth)}")
|
|
is Value -> value.toString()
|
|
is Split -> children.entries.joinToString("") { "\n${indent(depth)}${it.key}: ${it.value.asString(depth + 1)}" }
|
|
}
|
|
}
|
|
|
|
fun countLeafs(): Node<K, Int> {
|
|
return when (this) {
|
|
is Leaf -> Value(value.size)
|
|
is Value -> Value(1)
|
|
is Split -> Split(children.mapValues { it.value.countLeafs() })
|
|
}
|
|
}
|
|
|
|
private fun indent(depth: Int) = (0..(depth * 2)).joinToString("") { " " }
|
|
}
|
|
|
|
fun <K> Node<K, Int>.sort(): Node<K, Int> {
|
|
return when (this) {
|
|
is Node.Leaf -> this
|
|
is Node.Value -> this
|
|
is Node.Split -> Node.Split(this.children.sortByValues().mapValues { it.value.sort() })
|
|
}
|
|
}
|
|
|
|
private fun <K> Map<K, Node<K, Int>>.sortByValues(): Map<K, Node<K, Int>> {
|
|
val sortedEntries = entries.sortedBy {
|
|
when (val v = it.value) {
|
|
is Node.Value -> v.value
|
|
is Node.Leaf -> v.value.size
|
|
is Node.Split -> Int.MAX_VALUE
|
|
}
|
|
}
|
|
return LinkedHashMap<K, Node<K, Int>>().apply {
|
|
for ((key, value) in sortedEntries) {
|
|
this[key] = value
|
|
}
|
|
}
|
|
}
|