LeetCode in Kotlin

2509. Cycle Length Queries in a Tree

Hard

You are given an integer n. There is a complete binary tree with 2n - 1 nodes. The root of that tree is the node with the value 1, and every node with a value val in the range [1, 2n - 1 - 1] has two children where:

You are also given a 2D integer array queries of length m, where queries[i] = [ai, bi]. For each query, solve the following problem:

  1. Add an edge between the nodes with values ai and bi.
  2. Find the length of the cycle in the graph.
  3. Remove the added edge between nodes with values ai and bi.

Note that:

Return an array answer of length m where answer[i] is the answer to the ith query.

Example 1:

Input: n = 3, queries = [[5,3],[4,7],[2,3]]

Output: [4,5,3]

Explanation: The diagrams above show the tree of 23 - 1 nodes. Nodes colored in red describe the nodes in the cycle after adding the edge.

Example 2:

Input: n = 2, queries = [[1,2]]

Output: [2]

Explanation: The diagram above shows the tree of 22 - 1 nodes. Nodes colored in red describe the nodes in the cycle after adding the edge.

Constraints:

Solution

@Suppress("UNUSED_PARAMETER")
class Solution {
    fun cycleLengthQueries(n: Int, queries: Array<IntArray>): IntArray {
        val m = queries.size
        val res = IntArray(m)
        for (i in 0 until m) {
            var a = queries[i][0]
            var b = queries[i][1]
            var count = 1
            while (a != b) {
                if (a > b) {
                    a = a shr 1
                } else {
                    b = b shr 1
                }
                count++
            }
            res[i] = count
        }
        return res
    }
}