leetcode-go/written-exam/microsoft-2019-sim/2/main.go

package main

import (
	"fmt"
	"math"
)

// Get the distance of two nodes in binary tree

var finalDist int

func main() {
	distOf(1, 12, 12)
	fmt.Println(finalDist, 0, getDistof(12, 12))
	distOf(1, 12, 6)
	fmt.Println(finalDist, 1, getDistof(12, 6))
	distOf(1, 4, 5)
	fmt.Println(finalDist, 2, getDistof(4, 5))
	distOf(1, 1, 11)
	fmt.Println(finalDist, 3, getDistof(1, 11))
	distOf(1, 5, 6)
	fmt.Println(finalDist, 4, getDistof(5, 6))
	distOf(1, 6, 9)
	fmt.Println(finalDist, 5, getDistof(6, 9))
	distOf(1, 8, 15)
	fmt.Println(finalDist, 6, getDistof(8, 15))
	distOf(1, 21, 13)
	fmt.Println(finalDist, 7, getDistof(21, 13))
	distOf(1, 101, 5)
	fmt.Println(finalDist, 8, getDistof(101, 5))
	distOf(1, 11, 101)
	fmt.Println(finalDist, 9, getDistof(11, 101))
}

//           __1__
//         _/     \_
//        /         \
//      _2_         _3_
//     /   \       /   \
//    4     5     6     7
//   / \   / \   / \   / \
//  8   9 10 11 12 13 14 15
func distOf(root, p, q int) int {
	if p < root && q < root {
		return -1
	}
	left := distOf(root*2, p, q) + 1
	right := distOf(root*2+1, p, q) + 1
	if left != 0 && right != 0 {
		finalDist = left + right
		return left
	}
	if left != 0 {
		if root == p || root == q {
			finalDist = left
		}
		return left
	}
	if right != 0 {
		if root == p || root == q {
			finalDist = right
		}
		return right
	}
	if root == p || root == q {
		return 0
	}
	return -1
}

func getDistof(p, q int) (dist int) {
	lp := int(math.Log2(float64(p)))
	lq := int(math.Log2(float64(q)))
	if lp != lq {
		det := abs(lp - lq)
		dist += det
		if lp < lq {
			q >>= uint(det)
		} else {
			p >>= uint(det)
		}
	}
	for ; p != q; p, q = p/2, q/2 {
		dist += 2
	}
	return
}

func abs(x int) int {
	if x < 0 {
		return -x
	}
	return x
}