1377.frog-position-after-t-seconds.kt

/*
 * @lc app=leetcode id=1377 lang=kotlin
 *
 * [1377] Frog Position After T Seconds
 *
 * https://leetcode.com/problems/frog-position-after-t-seconds/description/
 *
 * algorithms
 * Hard (29.87%)
 * Likes:    26
 * Dislikes: 20
 * Total Accepted:    3K
 * Total Submissions: 10.1K
 * Testcase Example:  '7\n[[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]]\n2\n4'
 *
 * Given an undirected tree consisting of n vertices numbered from 1 to n. A
 * frog starts jumping from the 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] = [fromi, toi] means that exists an edge connecting directly the
 * vertices fromi and toi.
 * 
 * Return the probability that after t seconds the frog is on the vertex
 * target.
 * 
 * 
 * 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. 
 * 
 * 
 * Example 3:
 * 
 * 
 * Input: n = 7, edges = [[1,2],[1,3],[1,7],[2,4],[2,6],[3,5]], t = 20, target
 * = 6
 * Output: 0.16666666666666666
 * 
 * 
 * 
 * Constraints:
 * 
 * 
 * 1 <= n <= 100
 * edges.length == n-1
 * edges[i].length == 2
 * 1 <= edges[i][0], edges[i][1] <= n
 * 1 <= t <= 50
 * 1 <= target <= n
 * Answers within 10^-5 of the actual value will be accepted as correct.
 * 
 */

// @lc code=start
import java.util.LinkedList
class Solution {
    data class Node(val root: Int, val p: Double, val time: Int)

    fun frogPosition(n: Int, edges: Array<IntArray>, t: Int, target: Int): Double {
        var q = LinkedList<Node>()
        var vis = BooleanArray(n+1)
        var e = Array(n+1, {HashSet<Int>()})
        for (edge in edges) {
            e[edge[0]].add(edge[1])
            e[edge[1]].add(edge[0])
        }
        q.push(Node(1, 1.0, 0))
        vis[1] = true
        
        while (!q.isEmpty()) {
            var top = q.poll()
            if (top.time == t && top.root == target) {
                return top.p
            }

            var cnt = 0
            for (id in e[top.root])
                if (!vis[id]) cnt++
            if (top.time < t) {
                val p = top.p*1.0/cnt
                var ok = false
                for (id in e[top.root]) {
                    if (!vis[id]) {
                        vis[id] = true
                        q.offer(Node(id, p, top.time+1))
                        ok = true
                    }
                }
                if (!ok) q.offer(Node(top.root, top.p, top.time+1))
            }
        }
        return 0.0
    }
}
// @lc code=end