I had a really hard time getting my head around this problem, which contributed to major flailing on my part trying to get it sorted. I had several almost-right solutions, but I kept not getting the recursion right. I eventually sorted it out to “I need to check the subtree heights” as the right attack, but didn’t finally sort it until the 10th try or so.

This wold absolutely have been a fail in an interview. I’m getting much better with pointers – I had no trouble getting the recursions to work, and never had a nil pointer problem – but I just could not get the concept clear enough in my head, and that lead to my just not getting it right in the code.

After a lot of flailing I finally came up with a height-checking algorithm that was the best shot at it I could do, but I had to finally check with the solutions to spot my error, which was treating the left and right subtrees separately when they needed to be unified.

Final version after the fix did okay in the rankings: top 84% in speed and top 81% in memory, but the path to getting it was really bad. I need to spend a lot more time on tree-crawling.

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func isBalanced(root *TreeNode) bool {
    if root == nil || (root.Left == nil && root.Right == nil) {
        return true
    }
    _, which := heightCheck(root)
    return which
}

func heightCheck(root *TreeNode) (int, bool) {
    if root == nil {
        return 0, true
    }
 
    lH, lB := heightCheck(root.Left)
    rH, rB := heightCheck(root.Right)

    // I was way overcomplicatng this, tracking the
    // branches separately and not sending along a
    // give-up flag. Doing that helped a lot.
    if !lB || !rB  || abs(lH - rH) > 1 {
        return -1, false
    }

    return max(lH, rH) + 1, true
}

func abs(a int) int {
    if a > 0 {
        return a
    }
    return -a
}

func maxDepth(root *TreeNode) int {
    if root == nil {
        return 0
    }

    left := maxDepth(root.Left)
    right := maxDepth(root.Right)

    return max(left, right)+1
}


func max(a,b int) int {
    if a>b {
        return a
    }
    return b
}

Interestingly, if one adds a second recursion shortcut check for leaves, the memory usage gets much better: 99th percentile, but the execution time gets much worse: 54th percentile.