LeetCode/kt/most-profitable-path-in-a-tree.kt

class Solution {
    private data class Finder(val edges: Array<IntArray>, val bob: Int, val amount: IntArray) {
        private val n: Int = amount.size
        private val distances: IntArray = IntArray(n)
        private val tree: List<MutableList<Int>> = List(n) { mutableListOf() }

        init {
            edges.forEach { it ->
                val (u, v) = it[0] to it[1]

                tree[u].add(v)
                tree[v].add(u)
            }
        }

        private fun run(
            u: Int,
            parent: Int,
            time: Int,
        ): Int {
            var (maxIncome, maxChild) = 0 to Int.MIN_VALUE

            // Find the distances from Bob
            distances[u] =
                when {
                    u == bob -> 0
                    else -> n
                }

            tree[u].filter { it != parent }.forEach { v ->
                maxChild = listOf(maxChild, run(v, u, time + 1)).max()!!
                distances[u] = listOf(distances[u], distances[v] + 1).min()!!
            }

            when {
                // Alice reaches the node first
                distances[u] > time -> maxIncome += amount[u]
                // Alice and Bob reach the node at the same time
                distances[u] == time -> maxIncome += amount[u] / 2
            }

            return when {
                maxChild == Int.MIN_VALUE -> maxIncome
                else -> maxIncome + maxChild
            }
        }

        fun run(): Int = run(0, 0, 0)
    }

    fun mostProfitablePath(
        edges: Array<IntArray>,
        bob: Int,
        amount: IntArray,
    ): Int = Finder(edges, bob, amount).run()
}