Solving https://open.kattis.com/problems/moneymatters/statistics

moneymatters.kt

typealias Node = Int
typealias Edge = Int
typealias Color = Int

fun debugln(str: String) {
    // println(str)
    assert(!str.isEmpty())
}

fun main() {
    val nm = readLine()!!.split(" ").map{ it.toInt() }
    val vars = (0 until nm[0]).map { readLine()!!.toInt() }
    val edges = (0 until nm[1]).map { readLine()!!.split(" ").map { it.toInt() } }
    debugln("vars = $vars edges = $edges")
    
    // Map<Node -> List of Nodes>
    val graph = mutableMapOf<Node, MutableList<Node>>()
    edges.forEach {
        graph.getOrPut(it[0], { mutableListOf<Node>() }).add(it[1])
        graph.getOrPut(it[1], { mutableListOf<Node>() }).add(it[0])
    }
    debugln("g = $graph")

    val uncoloredNodes : MutableSet<Node> = (0 until vars.size).map { it }.toMutableSet()
    val nodeColors = mutableMapOf<Node, Color>()
    val coloring = mutableMapOf<Color, MutableList<Node>>()
    fun traverse(n: Node, c: Color) {
        if (!uncoloredNodes.contains(n)) {
            return
        }

        uncoloredNodes.remove(n)
        nodeColors[n] = c
        coloring.getOrPut(c, { mutableListOf<Node>() }).add(n)
        graph[n]!!.forEach { traverse(it, c) }
    }

    while (!uncoloredNodes.isEmpty()) {
        val start = uncoloredNodes.elementAt(0)
        val color = coloring.size
        traverse(start, color)
    }
    debugln("coloring = $coloring")

    if (coloring.values.any { it.map { vars[it] }.sum() != 0 }) {
        println("IMPOSSIBLE")
    } else {
        println("POSSIBLE")
    }
}