leetcode-go/oj/tencent-2019-summer/no6/ring.go

package main

import (
	"container/heap"
	"fmt"
)

type pair struct {
	_1 int
	_2 int
}

type maxHeap []pair

func (h maxHeap) Len() int           { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[j]._1 < h[i]._1 }
func (h maxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *maxHeap) Push(x interface{}) {
	*h = append(*h, x.(pair))
}

func (h *maxHeap) Pop() interface{} {
	l := len(*h)
	x := (*h)[l-1]
	*h = (*h)[:l-1]
	return x
}

func main() {
	var n, m int
	fmt.Scan(&n, &m)
	nums := make([]int, n)
	l, r := make([]int, n), make([]int, n)
	var h maxHeap
	for i := 0; i < n; i++ {
		fmt.Scan(&nums[i])
		heap.Push(&h, pair{nums[i], i})
		l[i] = (i + n - 1) % n
		r[i] = (i + 1) % n
	}
	sum := 0
	for i := 0; i < m; i++ {
		p := heap.Pop(&h).(pair)
		for l[p._2] == -1 {
			p = heap.Pop(&h).(pair)
		}
		sum += p._1
		left, right := l[p._2], r[p._2]
		l[p._2], r[p._2] = l[left], r[right]
		l[left], l[right] = -1, -1 // Mark element as used
		np := pair{nums[left] + nums[right] - p._1, p._2}
		heap.Push(&h, np)
	}
	fmt.Println(sum)
}