Created
May 23, 2022 23:21
-
-
Save imaginate/37c25fc2be4ffcfd01961be91d478af8 to your computer and use it in GitHub Desktop.
Leetcode 2281. Sum of Total Strength of Wizards - 8th Place Solution
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment