leetcode-go/leetcode/hard/363.max-sum-of-rectangle-no-larger-than-k.go

func maxSumSubmatrix(matrix [][]int, k int) (max int) {
	m := len(matrix)
	if m == 0 {
		return 0
	}
	n := len(matrix[0])
	if n == 0 {
		return 0
	}
	// Remember how to calculate the maximum sum of rectangles in a matrix:
	//
	// l, r              dp        l    r          dp                      l, r         dp
	//  -1   -4    0     -1        .   -4   .      -5                   .    .   .       .
	//   3   -3    4      3  --\   .   -3    .      0  --\   ...  --\   .    .    .      .
	//   4    8    7      4  --/   .    8     .    12  --/        --/   .    .     .     .
	//   6   20    1      6        .   20      .   26                   .    .      .    .
	//
	// Then, calculate the maximum sum of ranges in 1D array (dp).
	// Given the target k, the problem is : find the largest sum closest to k in dp.
	// Time compexity: O(n^2mlog(m)), n < m.
	max = -1 << 32
	if m < n { // Ensure n <= m
		m, n = n, m
		matrix = transpose(matrix)
	}
	for l := 0; l < n; l++ {
		dp := make([]int, m) // m rows
		for r := l; r < n; r++ {
			sum := []int{0}
			currSum := 0
			for i := 0; i < m; i++ {
				dp[i] += matrix[i][r]
				currSum += dp[i]
				idx := lower_bound(sum, currSum-k)
				if idx != len(sum) {
					max = maxInt(max, currSum-sum[idx])
				}
				insert(&sum, currSum)
			}
		}
	}
	return
}

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

func transpose(matrix [][]int) [][]int {
	m, n := len(matrix), len(matrix[0])
	res := make([][]int, n)
	for i := 0; i < n; i++ {
		res[i] = make([]int, m)
		for j := 0; j < m; j++ {
			res[i][j] = matrix[j][i]
		}
	}
	return res
}

func lower_bound(nums []int, val int) int {
	beg, end := 0, len(nums)-1
	for beg <= end {
		mid := beg + (end-beg)/2
		if nums[mid] < val {
			beg = mid + 1
		} else if val < nums[mid] {
			end = mid - 1
		} else {
			return mid
		}
	}
	return beg
}

func insert(nums *[]int, val int) { // nums is a set.
	i := lower_bound(*nums, val)
	if n := len(*nums); i == n {
		*nums = append(*nums, val)
	} else if (*nums)[i] != val {
		newNums := make([]int, n+1)
		copy(newNums, (*nums)[:i])
		newNums[i] = val
		copy(newNums[i+1:], (*nums)[i:])
		*nums = newNums
	}
}