BinaryIndexedTree, 数状数组

BinaryIndexedTree.kt

class BinaryIndexedTree(private var nums: IntArray) {
    var size: Int = nums.size
    private val _bit = IntArray(size + 1)

    init {
        for (i in 0 until size) init(i, nums[i])
    }

    private fun init(index: Int, value: Int) {
        var i = index + 1
        while (i <= size) {
            _bit[i] += value
            i += i and -i
        }
    }

    fun update(index: Int, value: Int) {
        val diff = value - nums[index]
        nums[index] = value
        init(index, diff)
    }

    fun sum(i: Int): Int {
        var i = i
        var sum = 0
        i++
        while (i > 0) {
            sum += _bit[i]
            i -= i and -i
        }
        return sum
    }

    fun sum(i: Int, j: Int): Int {
        return sum(j) - sum(i - 1)
    }
}