LeetCode in Kotlin

3187. Peaks in Array

Hard

A peak in an array arr is an element that is greater than its previous and next element in arr.

You are given an integer array nums and a 2D integer array queries.

You have to process queries of two types:

Return an array answer containing the results of the queries of the first type in order.

Notes:

Example 1:

Input: nums = [3,1,4,2,5], queries = [[2,3,4],[1,0,4]]

Output: [0]

Explanation:

First query: We change nums[3] to 4 and nums becomes [3,1,4,4,5].

Second query: The number of peaks in the [3,1,4,4,5] is 0.

Example 2:

Input: nums = [4,1,4,2,1,5], queries = [[2,2,4],[1,0,2],[1,0,4]]

Output: [0,1]

Explanation:

First query: nums[2] should become 4, but it is already set to 4.

Second query: The number of peaks in the [4,1,4] is 0.

Third query: The second 4 is a peak in the [4,1,4,2,1].

Constraints:

Solution

import kotlin.math.max

@Suppress("NAME_SHADOWING")
class Solution {
    fun countOfPeaks(nums: IntArray, queries: Array<IntArray>): List<Int> {
        val peaks = BooleanArray(nums.size)
        val binaryIndexedTree = IntArray(Integer.highestOneBit(peaks.size) * 2 + 1)
        for (i in 1 until peaks.size - 1) {
            if (nums[i] > max(nums[i - 1], nums[i + 1])) {
                peaks[i] = true
                update(binaryIndexedTree, i + 1, 1)
            }
        }
        val result: MutableList<Int> = ArrayList()
        for (query in queries) {
            if (query[0] == 1) {
                val leftIndex = query[1]
                val rightIndex = query[2]
                result.add(computeRangeSum(binaryIndexedTree, leftIndex + 2, rightIndex))
            } else {
                val index = query[1]
                val value = query[2]
                nums[index] = value
                for (i in -1..1) {
                    val affected = index + i
                    if (affected >= 1 && affected <= nums.size - 2) {
                        val peak =
                            nums[affected] > max(nums[affected - 1], nums[affected + 1])
                        if (peak != peaks[affected]) {
                            if (peak) {
                                update(binaryIndexedTree, affected + 1, 1)
                            } else {
                                update(binaryIndexedTree, affected + 1, -1)
                            }
                            peaks[affected] = peak
                        }
                    }
                }
            }
        }
        return result
    }

    private fun computeRangeSum(binaryIndexedTree: IntArray, beginIndex: Int, endIndex: Int): Int {
        return if (beginIndex <= endIndex) query(binaryIndexedTree, endIndex) - query(binaryIndexedTree, beginIndex - 1)
        else 0
    }

    private fun query(binaryIndexedTree: IntArray, index: Int): Int {
        var index = index
        var result = 0
        while (index != 0) {
            result += binaryIndexedTree[index]
            index -= index and -index
        }

        return result
    }

    private fun update(binaryIndexedTree: IntArray, index: Int, delta: Int) {
        var index = index
        while (index < binaryIndexedTree.size) {
            binaryIndexedTree[index] += delta
            index += index and -index
        }
    }
}