package main

import (
	"fmt"
	"math/rand"
	"sort"
	"time"
)

func main() {
	rand.Seed(time.Now().Unix())
	N := 100
	K := 10
	V := 100
	testTimes := 5000
	queryTimes := 100
	fmt.Println("测试开始")
	for i := 0; i < testTimes; i++ {
		n := rand.Intn(N) + 1
		k := rand.Intn(K) + 1
		arr := randomArray(n, V)
		st := NewSegmentTree(arr, k)
		right := NewRight(arr, k)
		for j := 0; j < queryTimes; j++ {
			a := rand.Intn(n)
			b := rand.Intn(n)
			l := GetMin(a, b)
			r := GetMax(a, b)
			ans1 := st.topKSum(l, r)
			ans2 := right.topKSum(l, r)
			if ans1 != ans2 {
				fmt.Println("出错了!")
				fmt.Println(ans1)
				fmt.Println(ans2)
				break
			}
		}
	}
	fmt.Println("测试结束")
}

func GetMax(a, b int) int {
	if a > b {
		return a
	} else {
		return b
	}
}

func GetMin(a, b int) int {
	if a < b {
		return a
	} else {
		return b
	}
}

type SegmentTree struct {
	n int
	k int
	// private int[] max;
	// private int[][] max = new int[][10];
	max   [][]int
	query [][]int
}

func NewSegmentTree(arr []int, K int) *SegmentTree {
	n := len(arr)
	k := K
	max := make([][]int, (n+1)<<2)
	query := make([][]int, (n+1)<<2)
	for i := 0; i < len(max); i++ {
		max[i] = make([]int, k)
		query[i] = make([]int, k)
	}
	ans := &SegmentTree{}
	ans.n = n
	ans.k = k
	ans.max = max
	ans.query = query
	for i := 0; i < n; i++ {
		ans.update(i, arr[i])
	}
	return ans
}

func (this *SegmentTree) topKSum(l int, r int) int {
	this.collect(l+1, r+1, 1, this.n, 1)
	sum := 0
	for _, num := range this.query[1] {
		sum += num
	}
	return sum
}

func (this *SegmentTree) update(i int, v int) {
	this.update0(i+1, i+1, v, 1, this.n, 1)
}

func (this *SegmentTree) update0(L int, R int, C int, l int, r int, rt int) {
	if L <= l && r <= R {
		this.max[rt][0] = C
		return
	}
	mid := (l + r) >> 1
	if L <= mid {
		this.update0(L, R, C, l, mid, rt<<1)
	}
	if R > mid {
		this.update0(L, R, C, mid+1, r, rt<<1|1)
	}
	this.merge(this.max[rt], this.max[rt<<1], this.max[rt<<1|1])
}

// father 要前k名
// left k名
// right k名
func (this *SegmentTree) merge(father []int, left []int, right []int) {
	for i, p1, p2 := 0, 0, 0; i < this.k; i++ {
		if left == nil || p1 == this.k {
			father[i] = right[p2]
			p2++
		} else if right == nil || p2 == this.k {
			father[i] = left[p1]
			p1++
		} else {
			if left[p1] >= right[p2] {
				father[i] = left[p1]
				p1++
			} else {
				father[i] = right[p2]
				p2++
			}
		}
	}
}

func (this *SegmentTree) collect(L int, R int, l int, r int, rt int) {
	if L <= l && r <= R {
		for i := 0; i < this.k; i++ {
			this.query[rt][i] = this.max[rt][i]
		}
	} else {
		mid := (l + r) >> 1
		leftUpdate := false
		rightUpdate := false
		if L <= mid {
			leftUpdate = true
			this.collect(L, R, l, mid, rt<<1)
		}
		if R > mid {
			rightUpdate = true
			this.collect(L, R, mid+1, r, rt<<1|1)
		}
		var left []int = nil
		if leftUpdate {
			left = this.query[rt<<1]
		}
		var right []int = nil
		if rightUpdate {
			right = this.query[rt<<1|1]
		}
		this.merge(this.query[rt], left, right)
	}
}

// // 暴力实现的结构
// // 为了验证
type Right struct {
	arr []int
	k   int
}

func NewRight(nums []int, K int) *Right {

	k := K
	arr := make([]int, len(nums))
	for i := 0; i < len(nums); i++ {
		arr[i] = nums[i]
	}
	return &Right{arr: arr, k: k}
}

func (this *Right) topKSum(l int, r int) int {
	heap := make([]int, 0)
	for i := l; i <= r; i++ {
		heap = append(heap, this.arr[i])
	}
	sum := 0
	for i := 0; i < this.k && len(heap) > 0; i++ {
		sort.Slice(heap, func(i, j int) bool {
			return heap[i] > heap[j]
		})
		sum += heap[0]
		heap = heap[1:]
	}
	return sum
}

// 为了验证
func randomArray(n int, v int) []int {
	ans := make([]int, n)
	for i := 0; i < n; i++ {
		ans[i] = rand.Intn(v) + 1
	}
	return ans
}