leetcode-go/medium/337.house-robber-iii.go

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
type node struct {
	root *TreeNode
	done bool
}

func rob(root *TreeNode) int {
	m := make(map[node]int)
	return maxInt(helper(root, false, &m), helper(root, true, &m))
}

func helper(root *TreeNode, done bool, m *map[node]int) int {
	if root == nil {
		return 0
	}
	if max, ok := (*m)[node{root, done}]; ok { // Memory search
		return max
	}
	var max int
	if done {
		max = helper(root.Left, false, m) + helper(root.Right, false, m)
	} else {
		max = maxInt(root.Val+helper(root.Left, true, m)+helper(root.Right, true, m), helper(root.Left, false, m)+helper(root.Right, false, m))
	}
	(*m)[node{root, done}] = max
	return max
}

func maxInt(x, y int) int {
	if x < y {
		return y
	}
	return x
}