Leetcode 2281. Sum of Total Strength of Wizards - 8th Place Solution

solution.kt

import java.util.*

const val MOD = 1000000007L

class Solution {
    fun totalStrength(strength: IntArray): Int {
        val n = strength.size
        val left = IntArray(n) { -1 }
        val right = IntArray(n) { -1 }
        val stack = Stack<Int>()
        for (j in 0 until n) {
            while (stack.isNotEmpty() && strength[j] <= strength[stack.peek()]) {
                left[j] = stack.pop()
            }
            if (stack.isNotEmpty()) {
                right[stack.peek()] = j
            }
            stack.push(j)
        }
        val sums = LongArray(n + 1)
        for (j in 0 until n) {
            sums[j + 1] = sums[j] + strength[j].toLong()
            sums[j + 1] %= MOD
        }
        val sums2 = LongArray(n + 1)
        for (j in 0 until n) {
            sums2[j + 1] = sums2[j] + sums[j + 1]
            sums2[j + 1] %= MOD
        }
        val sums3 = LongArray(n + 1)
        for (j in n - 1 downTo 0) {
            sums3[j] = sums3[j + 1] + (sums[n] - sums[j])
            sums3[j] %= MOD
        }
        var answer = 0L
        fun recur(from: Int, mid: Int, to: Int) {
            if (mid != -1) {
                var leftAMT = sums3[from] - sums3[mid] - ((mid - from).toLong() * (sums[n] - sums[mid]))
                leftAMT %= MOD
                var rightAMT = sums2[to + 1] - sums2[mid + 1] - ((to - mid).toLong() * sums[mid + 1])
                rightAMT %= MOD
                var here = 0L
                here += (to - mid + 1).toLong() * leftAMT
                here %= MOD
                here += (mid - from + 1).toLong() * rightAMT
                here %= MOD
                here += (((to - mid + 1).toLong() * (mid - from + 1).toLong()) % MOD) * strength[mid].toLong()
                here %= MOD
                answer += here * strength[mid].toLong()
                answer %= MOD
                recur(from, left[mid], mid - 1)
                recur(mid + 1, right[mid], to)
            }
        }
        while (stack.size > 1) {
            stack.pop()
        }
        recur(0, stack.pop(), n - 1)
        answer += MOD
        answer %= MOD
        return answer.toInt()
    }
}