huffman.kt

sealed class HuffmanNode {
    abstract val codeSize: Int

    data class CodeInfo(val code: Char, override val codeSize: Int) : HuffmanNode()
    data class TreeNode(val left: HuffmanNode, val right: HuffmanNode, override val codeSize: Int) : HuffmanNode()
} 

fun main() {
    val text = "HELLO"
    val result = huffman(text)
    println(result)
}

fun huffman(text: String): HuffmanNode {
    val huffmanNodes = text.fold(mutableMapOf<Char, Int>()) { acc, c ->
        val n = acc[c] ?: 0
        acc[c] = n + 1
        acc
    }.asSequence()
        .fold(PriorityQueue<HuffmanNode>(
            Comparator { o1, o2 ->
                o1.codeSize.compareTo(o2.codeSize)
            })
        ) { acc, entry ->
            acc.add(HuffmanNode.CodeInfo(entry.key, entry.value))
            acc
        }
    return when (huffmanNodes.size) {
        1 -> huffmanNodes.first()
        else -> {
            while (huffmanNodes.size != 1) {
                val left = huffmanNodes.poll()
                val right = huffmanNodes.poll()
                val node = HuffmanNode.TreeNode(left, right, left.codeSize + right.codeSize)
                huffmanNodes.add(node)
            }
            huffmanNodes.poll()
        }
    }
}