LeetCode in Kotlin

1377. Frog Position After T Seconds

Hard

Given an undirected tree consisting of n vertices numbered from 1 to n. A frog starts jumping from vertex 1. In one second, the frog jumps from its current vertex to another unvisited vertex if they are directly connected. The frog can not jump back to a visited vertex. In case the frog can jump to several vertices, it jumps randomly to one of them with the same probability. Otherwise, when the frog can not jump to any unvisited vertex, it jumps forever on the same vertex.

The edges of the undirected tree are given in the array edges, where edges[i] = [ai, bi] means that exists an edge connecting the vertices ai and bi.

Return the probability that after t seconds the frog is on the vertex target. Answers within 10-5 of the actual answer will be accepted.

Example 1:

Input: n = 7, edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]], t = 2, target = 4

Output: 0.16666666666666666

Explanation: The figure above shows the given graph. The frog starts at vertex 1, jumping with 1/3 probability to the vertex 2 after second 1 and then jumping with 1/2 probability to vertex 4 after second 2. Thus the probability for the frog is on the vertex 4 after 2 seconds is 1/3 * 1/2 = 1/6 = 0.16666666666666666.

Example 2:

Input: n = 7, edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]], t = 1, target = 7

Output: 0.3333333333333333

Explanation: The figure above shows the given graph. The frog starts at vertex 1, jumping with 1/3 = 0.3333333333333333 probability to the vertex 7 after second 1.

Constraints:

Solution

import java.util.LinkedList
import java.util.Queue

@Suppress("NAME_SHADOWING")
class Solution {
    fun frogPosition(n: Int, edges: Array<IntArray>, t: Int, target: Int): Double {
        var t = t
        val graph: Array<MutableList<Int>?> = arrayOfNulls(n)
        for (i in 0 until n) {
            graph[i] = ArrayList()
        }
        for (edge in edges) {
            graph[edge[0] - 1]?.add(edge[1] - 1)
            graph[edge[1] - 1]?.add(edge[0] - 1)
        }
        val visited = BooleanArray(n)
        visited[0] = true
        val probabilities = DoubleArray(n)
        probabilities[0] = 1.0
        val queue: Queue<Int> = LinkedList()
        queue.offer(0)
        while (queue.isNotEmpty() && t-- > 0) {
            for (i in queue.size downTo 1) {
                val vertex = queue.poll()
                var nextVerticesCount = 0
                for (next in graph[vertex]!!) {
                    if (!visited[next]) {
                        nextVerticesCount++
                    }
                }
                for (next in graph[vertex]!!) {
                    if (!visited[next] && nextVerticesCount > 0) {
                        visited[next] = true
                        queue.offer(next)
                        probabilities[next] = probabilities[vertex] / nextVerticesCount
                    }
                }
                if (nextVerticesCount > 0) {
                    probabilities[vertex] = 0.0
                }
            }
        }
        return probabilities[target - 1]
    }
}