LeetCode in Kotlin

3558. Number of Ways to Assign Edge Weights I

Medium

There is an undirected tree with n nodes labeled from 1 to n, rooted at node 1. The tree is represented by a 2D integer array edges of length n - 1, where edges[i] = [ui, vi] indicates that there is an edge between nodes ui and vi.

Initially, all edges have a weight of 0. You must assign each edge a weight of either 1 or 2.

The cost of a path between any two nodes u and v is the total weight of all edges in the path connecting them.

Select any one node x at the maximum depth. Return the number of ways to assign edge weights in the path from node 1 to x such that its total cost is odd.

Since the answer may be large, return it modulo 109 + 7.

Note: Ignore all edges not in the path from node 1 to x.

Example 1:

Input: edges = [[1,2]]

Output: 1

Explanation:

• The path from Node 1 to Node 2 consists of one edge (1 → 2).
• Assigning weight 1 makes the cost odd, while 2 makes it even. Thus, the number of valid assignments is 1.

Example 2:

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

Output: 2

Explanation:

• The maximum depth is 2, with nodes 4 and 5 at the same depth. Either node can be selected for processing.
• For example, the path from Node 1 to Node 4 consists of two edges (1 → 3 and 3 → 4).
• Assigning weights (1,2) or (2,1) results in an odd cost. Thus, the number of valid assignments is 2.

Constraints:

• 2 <= n <= 105
• edges.length == n - 1
• edges[i] == [ui, vi]
• 1 <= ui, vi <= n
• edges represents a valid tree.

Solution

class Solution {
    fun assignEdgeWeights(edges: Array<IntArray>): Int {
        if (pow2[0] == 0L) {
            pow2[0] = 1
            for (i in 1..<pow2.size) {
                pow2[i] = (pow2[i - 1] shl 1) % mod
            }
        }
        val n = edges.size + 1
        val adj = IntArray(n + 1)
        val degrees = IntArray(n + 1)
        for (edge in edges) {
            val u = edge[0]
            val v = edge[1]
            adj[u] += v
            adj[v] += u
            degrees[u]++
            degrees[v]++
        }
        val que = IntArray(n)
        var write = 0
        var read = 0
        for (i in 2..n) {
            if (degrees[i] == 1) {
                que[write++] = i
            }
        }
        var distance = 0
        while (read < write) {
            distance++
            var size = write - read
            while (size-- > 0) {
                val v = que[read++]
                val u = adj[v]
                adj[u] -= v
                if (--degrees[u] == 1 && u != 1) {
                    que[write++] = u
                }
            }
        }
        return pow2[distance - 1].toInt()
    }

    companion object {
        private const val mod = 1e9.toInt() + 7
        private val pow2 = LongArray(100001)
    }
}

This site is open source. Improve this page.