LeetCode/kt/the-k-th-lexicographical-string-of-all-happy-strings-of-length-n.kt

class Solution {
    companion object {
        private val NEXT_SMALLEST: Map<Char, Char> =
            mapOf(
                'a' to 'b',
                'b' to 'a',
                'c' to 'a',
            )
        private val NEXT_GREATEST: Map<Char, Char> =
            mapOf(
                'a' to 'c',
                'b' to 'c',
                'c' to 'b',
            )
    }

    private fun variations(n: Int): Int = 1.shl(n - 1)

    private fun maxHappyStrings(n: Int): Int = 3 * variations(n)

    fun getHappyString(
        n: Int,
        k: Int,
    ): String {
        var k = k

        // Find how many strings can be constructed
        val total = maxHappyStrings(n)
        if (k > total) {
            return ""
        }

        // Initialize the resulting string
        val result = StringBuilder(n)
        (1..n).forEach {
            result.append('a')
        }

        val startA = 1
        val startB = startA + variations(n)
        val startC = startB + variations(n)

        when {
            k < startB -> {
                result.setCharAt(0, 'a')
                k -= startA
            }
            k < startC -> {
                result.setCharAt(0, 'b')
                k -= startB
            }
            else -> {
                result.setCharAt(0, 'c')
                k -= startC
            }
        }

        (1..n - 1).forEach { index ->
            val mid: Int = 1.shl(n - index - 1)

            when {
                k < mid -> result.setCharAt(index, NEXT_SMALLEST[result[index - 1]]!!)
                else -> {
                    result.setCharAt(index, NEXT_GREATEST[result[index - 1]]!!)
                    k -= mid
                }
            }
        }

        return result.toString()
    }
}