leetcode-go/leetcode/hard/632.smallest-range.go

type array [][]int

func (a array) Len() int           { return len(a) }
func (a array) Less(i, j int) bool { return a[i][0] < a[j][0] }
func (a array) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }

func smallestRange(nums [][]int) []int {
	var a array
	m := make(map[int]int)
	for i := range nums {
		for j := range nums[i] {
			if num := nums[i][j]; j == 0 || num != nums[i][j-1] {
				a = append(a, []int{num, i})
			}
		}
	}
	sort.Sort(a)
	l, r, k, n := 0, 0, len(nums), len(a)
	min := math.MaxInt32
	res := make([]int, 2)
	for r < n {
		for r < n && len(m) != k {
			for {
				m[a[r][1]]++
				r++
				if n-1 < r || a[r-1][0] != a[r][0] {
					break
				}
			}
		}
		if len(m) == k && a[r-1][0]-a[l][0] < min {
			res[0], res[1] = a[l][0], a[r-1][0]
			min = a[r-1][0] - a[l][0]
		}
		for l < r && len(m) == k {
			for {
				m[a[l][1]]--
				if m[a[l][1]] == 0 {
					delete(m, a[l][1])
				}
				l++
				if len(m) == k && a[r-1][0]-a[l][0] < min {
					res[0], res[1] = a[l][0], a[r-1][0]
					min = a[r-1][0] - a[l][0]
				}
				if l == r || a[l-1][0] != a[l][0] {
					break
				}
			}
		}
	}
	return res
}