LeetCode in Kotlin

1786. Number of Restricted Paths From First to Last Node

Medium

There is an undirected weighted connected graph. You are given a positive integer n which denotes that the graph has n nodes labeled from 1 to n, and an array edges where each edges[i] = [ui, vi, weighti] denotes that there is an edge between nodes ui and vi with weight equal to weighti.

A path from node start to node end is a sequence of nodes [z0, z1, z2, ..., zk] such that z0 = start and zk = end and there is an edge between zi and zi+1 where 0 <= i <= k-1.

The distance of a path is the sum of the weights on the edges of the path. Let distanceToLastNode(x) denote the shortest distance of a path between node n and node x. A restricted path is a path that also satisfies that distanceToLastNode(zi) > distanceToLastNode(zi+1) where 0 <= i <= k-1.

Return the number of restricted paths from node 1 to node n. Since that number may be too large, return it modulo 109 + 7.

Example 1:

Input: n = 5, edges = [[1,2,3],[1,3,3],[2,3,1],[1,4,2],[5,2,2],[3,5,1],[5,4,10]]

Output: 3

Explanation: Each circle contains the node number in black and its distanceToLastNode value in blue. The three restricted paths are:

1) 1 –> 2 –> 5

2) 1 –> 2 –> 3 –> 5

3) 1 –> 3 –> 5

Example 2:

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

Output: 1

Explanation: Each circle contains the node number in black and its distanceToLastNode value in blue. The only restricted path is 1 –> 3 –> 7.

Constraints:

Solution

import java.util.PriorityQueue

class Solution {
    private class Pair(var v: Int, var cwt: Int) : Comparable<Pair> {
        override fun compareTo(other: Pair): Int {
            return cwt - other.cwt
        }
    }

    private class Edge(var v: Int, var wt: Int)

    private lateinit var dtl: IntArray
    private lateinit var dp: IntArray

    fun countRestrictedPaths(n: Int, edges: Array<IntArray>): Int {
        val graph = buildGraph(n, edges)
        val pq = PriorityQueue<Pair>()
        val vis = BooleanArray(n + 1)
        dtl = IntArray(n + 1)
        pq.add(Pair(n, 0))
        while (pq.isNotEmpty()) {
            val rem = pq.remove()
            if (vis[rem.v]) {
                continue
            }
            dtl[rem.v] = rem.cwt
            vis[rem.v] = true
            for (edge in graph[rem.v]) {
                if (!vis[edge.v]) {
                    pq.add(Pair(edge.v, rem.cwt + edge.wt))
                }
            }
        }
        dp = IntArray(n + 1)
        return dfs(graph, 1, BooleanArray(n + 1), n)
    }

    private fun dfs(graph: List<MutableList<Edge>>, vtx: Int, vis: BooleanArray, n: Int): Int {
        if (vtx == n) {
            return 1
        }
        var ans: Long = 0
        vis[vtx] = true
        for (edge in graph[vtx]) {
            if (!vis[edge.v] && dtl[edge.v] < dtl[vtx]) {
                val x = dfs(graph, edge.v, vis, n)
                ans = (ans + x) % M
            } else if (dtl[edge.v] < dtl[vtx] && vis[edge.v]) {
                ans = (ans + dp[edge.v]) % M
            }
        }
        dp[vtx] = ans.toInt()
        return ans.toInt()
    }

    private fun buildGraph(n: Int, edges: Array<IntArray>): List<MutableList<Edge>> {
        val graph: MutableList<MutableList<Edge>> = ArrayList()
        for (i in 0..n) {
            graph.add(ArrayList())
        }
        for (edge in edges) {
            val u = edge[0]
            val v = edge[1]
            val wt = edge[2]
            graph[u].add(Edge(v, wt))
            graph[v].add(Edge(u, wt))
        }
        return graph
    }

    companion object {
        private const val M = 1000000007
    }
}