FB question

balance.kt

import java.util.*

fun main(args : Array<String>) {
    val numBalances = readLine()!!.toInt()
    val balances = (1 .. numBalances).map {
        val leftLine = readLine()!!.split(' ')
        val rightLine = readLine()!!.split(' ')

        Balance(leftLine[0].toInt(), if (leftLine.size == 2) leftLine[1].toInt() else -1,
                rightLine[0].toInt(), if (rightLine.size == 2) rightLine[1].toInt() else -1)
    }

    // need to do a BFS and record the visiting order
    val bfsQueue : Queue<Int> = LinkedList()
    val balanceOrder : Stack<Int> = Stack()

    // assuming 0 is always the beginning
    bfsQueue.add(0)
    while(bfsQueue.isNotEmpty()) {
        val i = bfsQueue.poll()!!
        balanceOrder.push(i)
        if (balances[i].scaleL != -1) {
            bfsQueue.add(balances[i].scaleL)
        }

        if (balances[i].scaleR != -1) {
            bfsQueue.add(balances[i].scaleR)
        }
    }

    val answers = Array(numBalances, { i -> 0 })

    while(balanceOrder.isNotEmpty()) {
        val i = balanceOrder.pop()!!
        val b = balances[i]
        val left = b.weightL + getBalanceWeight(balances, b.scaleL) + ( if (b.scaleL != -1) Math.abs(answers[b.scaleL]) else 0 )
        val right = b.weightR + getBalanceWeight(balances, b.scaleR) + ( if (b.scaleR != -1) Math.abs(answers[b.scaleR]) else 0 )
        answers[i] = left - right
    }

    for ((index, answer) in answers.withIndex()) {
        println("$index: ${if(answer > 0) "0 $answer" else "${-answer} 0"}")
    }
}

fun getBalanceWeight(balances : List<Balance>, balanceId : Int) : Int {
    return when(balanceId) {
        -1 -> 0
        else -> {
            val b = balances[balanceId]
            return 10 +
                b.weightL + getBalanceWeight(balances, b.scaleL) +
                b.weightR + getBalanceWeight(balances, b.scaleR)
        }
    }
}

data class Balance (val weightL : Int, val scaleL: Int,
                    val weightR : Int, val scaleR: Int)