LeetCode 对角线遍历II

/**
LeetCode 对角线遍历 II
对于同一个对角线上，i与j相加的值是相同的，而且i越大，越早输出
由此，可以通过设置一个比较大的数值为一个位置的权重，同一个对角线的权重相同，而i也是影响输出时机
以此基础，可以通过排序，获的输出时机
*/
func findDiagonalOrder(nums [][]int) []int {
	cnt := len(nums)
	// 权重初始化
	maxW := 1000000000
	dias := make([]*Dia, 0)
	for i:=0; i<cnt; i++ {
		jCnt := len(nums[i])
		for j:=0; j<jCnt; j++ {
			dia := &Dia{
				w: maxW-i-j, //同一个对角线上的权重相同
				p: i,        //同一个对角线上的i越大，越早输出
				val: nums[i][j],
			}
			dias = append(dias, dia)
		}
	}
	//根据权重与i进行排序
	sort.Sort(DiaWapper{Dias:dias, By: func(p, q *Dia) bool {
		if p.w == q.w {
			return q.p < p.p
		}
		return q.w < p.w
	}})
	diaLen := len(dias)
	ans := make([]int, 0)
	for i:=0; i<diaLen; i++ {
		ans = append(ans, dias[i].val)
	}
	return ans
}

type Dia struct {
	w, p int //权重与i
	val int  //nums[i][j]
}
type DiaWapper struct {
	Dias []*Dia
	By func(p, q *Dia) bool
}
func (dia DiaWapper) Swap(i, j int) {
	dia.Dias[i], dia.Dias[j] = dia.Dias[j], dia.Dias[i]
}
func (dia DiaWapper) Less(i, j int) bool {
	return dia.By(dia.Dias[i], dia.Dias[j])
}
func (dia DiaWapper) Len() int {
	return len(dia.Dias)
}