LeetCode in Kotlin

3067. Count Pairs of Connectable Servers in a Weighted Tree Network

Medium

You are given an unrooted weighted tree with n vertices representing servers numbered from 0 to n - 1, an array edges where edges[i] = [a_i, b_i, weight_i] represents a bidirectional edge between vertices a_i and b_i of weight weight_i. You are also given an integer signalSpeed.

Two servers a and b are connectable through a server c if:

a < b, a != c and b != c.
The distance from c to a is divisible by signalSpeed.
The distance from c to b is divisible by signalSpeed.
The path from c to b and the path from c to a do not share any edges.

Return an integer array count of length n where count[i] is the number of server pairs that are connectable through the server i.

Example 1:

Input: edges = [[0,1,1],[1,2,5],[2,3,13],[3,4,9],[4,5,2]], signalSpeed = 1

Output: [0,4,6,6,4,0]

Explanation: Since signalSpeed is 1, count[c] is equal to the number of pairs of paths that start at c and do not share any edges.

In the case of the given path graph, count[c] is equal to the number of servers to the left of c multiplied by the servers to the right of c.

Example 2:

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

Output: [2,0,0,0,0,0,2]

Explanation: Through server 0, there are 2 pairs of connectable servers: (4, 5) and (4, 6).

Through server 6, there are 2 pairs of connectable servers: (4, 5) and (0, 5).

It can be shown that no two servers are connectable through servers other than 0 and 6.

Constraints:

2 <= n <= 1000
edges.length == n - 1
edges[i].length == 3
0 <= a_i, b_i < n
edges[i] = [a_i, b_i, weight_i]
1 <= weight_i <= 10⁶
1 <= signalSpeed <= 10⁶
The input is generated such that edges represents a valid tree.

Solution

@Suppress("NAME_SHADOWING")
class Solution {
    private lateinit var adj: Array<ArrayList<Int>>

    fun countPairsOfConnectableServers(edges: Array<IntArray>, signalSpeed: Int): IntArray {
        val n = edges.size + 1
        adj = Array(n) { ArrayList() }
        for (i in 0 until n) {
            adj[i] = ArrayList()
        }
        for (edge in edges) {
            val u = edge[0]
            val v = edge[1]
            val w = edge[2]
            adj[u].add(v)
            adj[v].add(u)
            adj[u].add(w)
            adj[v].add(w)
        }
        val res = IntArray(n)
        for (i in 0 until n) {
            if (adj[i].size > 2) {
                val al = ArrayList<Int>()
                var j = 0
                while (j < adj[i].size) {
                    val cnt = IntArray(1)
                    dfs(adj[i][j], i, adj[i][j + 1], cnt, signalSpeed)
                    al.add(cnt[0])
                    j += 2
                }
                var sum = 0
                for (j in al) {
                    res[i] += (sum * j)
                    sum += j
                }
            }
        }
        return res
    }

    fun dfs(node: Int, par: Int, sum: Int, cnt: IntArray, ss: Int) {
        if (sum % ss == 0) {
            cnt[0]++
        }
        var i = 0
        while (i < adj[node].size) {
            val child = adj[node][i]
            if (child != par) {
                dfs(child, node, sum + adj[node][i + 1], cnt, ss)
            }
            i += 2
        }
    }
}

This site is open source. Improve this page.