leetcode-go/leetcode/hard/546.remove-boxes.go

const MAX_N int = 100

var dp map[int]int
var length []int

func removeBoxes(boxes []int) int {
	n := len(boxes)
	dp = make(map[int]int)
	length = make([]int, n)
	for i := 1; i < n; i++ {
		if boxes[i] == boxes[i-1] {
			length[i] = length[i-1] + 1
		}
	}
	return dfs(boxes, 0, len(boxes)-1, 0)
}

func dfs(boxes []int, l, r, k int) int {
	// dp[l][r][k] = dp[l][r-1][0] + (k+1)^2, drop box[r]
	//             = dp[l][p][k+1] + dp[p+1][r-1][0],
	// box[p] == box[j], remove boxes at [p+1, r-1] and attach
	// box[j] to box[p]
	if r < l {
		return 0
	}
	k += length[r]
	r -= length[r]
	key := ((l*MAX_N)+r)*MAX_N + k
	if val, ok := dp[key]; ok {
		return val
	}
	res := dfs(boxes, l, r-1, 0) + (k+1)*(k+1)
	for i := l; i < r; i++ {
		if boxes[i] == boxes[r] {
			res = maxInt(res, dfs(boxes, l, i, k+1)+dfs(boxes, i+1, r-1, 0))
		}
	}
	dp[key] = res
	return res
}

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