[baekjoon] #1275: 커피숍2

problem link

• simple segment tree problem
• only update in-range index of segment tree to reduce time complexity

lateinit var segTree: Array<Long>
lateinit var nums: MutableList<Long>

fun main() {
    val br = System.`in`.bufferedReader()
    val bw = System.out.bufferedWriter()

    val (n, q) = br.readLine().split(" ").map { it.toInt() }
    nums = br.readLine().split(" ").map { it.toLong() }.toMutableList()

    val segTreeSize = 4 * n
    segTree = Array(segTreeSize) { 0L }
    constructSegTree(0, n - 1, 0)
    
    for (i in 0 until q) {
        var (x, y, a, b) = br.readLine().split(" ").map { it.toInt() }
        // make x < y
        if (x > y) {
            val temp = x
            x = y
            y = temp
        }
        // all indexes start with 0
        bw.write("${getSegTreeValue(0, n - 1, x - 1, y - 1, 0)}\n")
        updateSegTree(0, n - 1, a - 1, b.toLong(), 0)
    }

    br.close()
    bw.close()
}

fun constructSegTree(low: Int, high: Int, pos: Int) {
    if (low == high) {
        segTree[pos] = nums[low]
        return
    }

    val mid = (low + high) / 2
    constructSegTree(low, mid, 2 * pos + 1)
    constructSegTree(mid + 1, high, 2 * pos + 2)
    segTree[pos] = segTree[2 * pos + 1] + segTree[2 * pos + 2]
}

fun getSegTreeValue(low: Int, high: Int, qLow: Int, qHigh: Int, pos: Int): Long {
    // total overlap
    if (qLow <= low && qHigh >= high) return segTree[pos]
    // no overlap
    if (qLow > high || qHigh < low) return 0

    // partial overlap
    val mid = (low + high) / 2
    return getSegTreeValue(low, mid, qLow, qHigh, 2 * pos + 1) +
            getSegTreeValue(mid + 1, high, qLow, qHigh, 2 * pos + 2)
}

fun updateSegTree(low: Int, high: Int, index: Int, value: Long, pos: Int) {
    // leaf node
    if (low == high) {
        // update value
        if (low == index) segTree[pos] = value
        return
    }
    // out of range
    if (index !in low..high) return

    // update in range values
    val mid = (low + high) / 2
    updateSegTree(low, mid, index, value, 2 * pos + 1)
    updateSegTree(mid + 1, high, index, value, 2 * pos + 2)
    segTree[pos] = segTree[2 * pos + 1] + segTree[2 * pos + 2]
}

Time Complexity

O(N)