Scala 树状数组BIT模板

Problem: 6404. 将数组清空

目录

• Code

Code

object Solution {
  def countOperationsToEmptyArray(nums: Array[Int]): Long = {
    val n = nums.length
    val id = Array.tabulate(n)(i => i)
    val sortedId = id.sortWith((i, j) => nums(i) - nums(j) < 0)

    var ans = n.toLong  //先把n计入答案
    val t = new BIT(n + 1)  //下标从1开始
    var pre = 1  //上一个最小值的位置，初始为1

    for (k <- sortedId.indices) {
      val i = sortedId(k) + 1 //下标从1开始
      if (i >= pre)  //从pre移动到i，跳过已经删除的数
        ans += i - pre - t.query(pre, i)
      else  //从pre移动到n，再从1移动到i，跳过已经删除的数
        ans += n - pre + i - k + t.query(i, pre - 1)
      t.inc(i) //删除i
      pre = i
    }
    ans
  }
}

//树状数组模板
class BIT(n: Int) {
  private val tree = new Array[Int](n)

  //将下标i位置的数加一
  def inc(i: Int): Unit = {
    var index = i
    while (index < tree.length) {
      tree(index) += 1
      index += index & -index
    }
  }

 //返回闭区间[1,i]的元素和
  def sum(x: Int): Int = {
    var res = 0
    var index = x
    while (index > 0) {
      res += tree(index)
      index &= index - 1
    }
    res
  }

  //返回闭区间[left,right]的元素和
  def query(left: Int, right: Int): Int = {
    sum(right) - sum(left - 1)
  }
}