Untitled

mail@pastecode.io avatar
unknown
kotlin
a year ago
1.4 kB
8
Indexable
Never
class NumArray(val nums: IntArray) {

    private val tree = IntArray(2 * nums.size)

    init {
        build(0, 0, nums.size - 1)
    }

    private fun build(v: Int, tl: Int, tr: Int) {
        if (tl > tr) return

        if (tl == tr) {
            tree[v] = nums[tl]
        } else {
            val mid = tl + (tr - tl) / 2
            build(2 * v + 1, tl, mid)
            build(2 * v + 2, mid + 1, tr)
            tree[v] = tree[2 * v + 1] + tree[2 * v + 2]
        }
    }

    fun update(index: Int, `val`: Int) {
        fun set(v: Int, tl: Int, tr: Int) {
            if (tl == tr) { 
                tree[v] = `val`
                return
            }

            val mid = (tl + tr) / 2
            if (index <= mid) 
                set(2 * v + 1, tl, mid)
            else
                set(2 * v + 2, mid + 1, tr)
            tree[v] = tree[2 * v + 1] + tree[2 * v + 2]
        }
        set(0, 0, nums.size - 1)
    }

    fun sumRange(left: Int, right: Int): Int {
        fun sum(v: Int, tl: Int, tr: Int): Int =
            if (left > tr || right < tl) 
                0
            else if (tl >= left && tr <= right)
                tree[v]
            else {
                val mid = (tl + tr) / 2
                sum(2 * v + 1, tl, mid) + sum(2 * v + 2, mid + 1, tr)
            }
        return sum(0, 0, nums.size - 1)
    }
}