2022-12-16：给你一个长度为n的数组，并询问q次每次询问区间[l,r]之间是否存在小于等于k个数的和大于等于x每条查询返

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
}