fun journeyToMoon(n: Int, astronaut: Array<Array<Int>>): Long {
    val nodes = (0..<n).map {Node(it)}
    val visited = mutableSetOf<Node>()
    val groups = mutableListOf<Set<Node>>()
    for (link in astronaut) {
        nodes[link[0]].connections.add(nodes[link[1]])
        nodes[link[1]].connections.add(nodes[link[0]])
    }
    for (node in nodes) {
        if (visited.contains(node)) continue
        val group = findConnections(node)
        groups.add(group)
        visited.addAll(group)
    }
    val sizes = groups.map {it.size.toLong()}
    val allPairs = n.toLong()*(n-1)/2
    val sameCountryPairs = sizes.map {it*(it-1)/2}
        .sum()
    return allPairs - sameCountryPairs
}

data class Node(val id: Int) {
    val connections = mutableSetOf<Node>()
}

fun findConnections(start: Node): Set<Node> {
    val visited = mutableSetOf<Node>()
    val queue = ArrayDeque<Node>()
    queue.add(start)
    while (!queue.isEmpty()) {
        val curr = queue.removeFirst()
        visited.add(curr)
        queue.addAll(curr.connections.filter {!visited.contains(it)})
    }
    return visited
}