go: add «2818. Apply Operations to Maximize Score»

URL:	https://leetcode.com/problems/apply-operations-to-maximize-score/
Signed-off-by: Matej Focko <me@mfocko.xyz>

go/apply-operations-to-maximize-score.go

@@ -0,0 +1,105 @@
+package main
+
+import (
+	"cmp"
+
+	pq "github.com/emirpasic/gods/v2/queues/priorityqueue"
+	arrst "github.com/emirpasic/gods/v2/stacks/arraystack"
+)
+
+func maximumScore(nums []int, k int) int {
+	MOD := int64(1000000007)
+
+	subCmp := func(a, b int) int {
+		return -cmp.Compare(nums[a], nums[b])
+	}
+
+	power := func(base, exp int64) int64 {
+		res := int64(1)
+
+		for exp > 0 {
+			if exp%2 == 1 {
+				res = (res * base) % MOD
+			}
+
+			base = (base * base) % MOD
+			exp >>= 1
+		}
+
+		return res
+	}
+
+	calculatePrimeScores := func() []int {
+		scores := make([]int, len(nums))
+		for i, num := range nums {
+			for factor := 2; factor*factor <= num; factor++ {
+				if num%factor != 0 {
+					continue
+				}
+
+				scores[i]++
+				for num%factor == 0 {
+					num /= factor
+				}
+			}
+
+			if num >= 2 {
+				scores[i]++
+			}
+		}
+
+		return scores
+	}
+
+	calculateSubarrays := func(scores []int) []int64 {
+		prev, next := make([]int, len(nums)), make([]int, len(nums))
+		for i := range nums {
+			prev[i], next[i] = -1, len(nums)
+		}
+
+		st := arrst.New[int]()
+		for i, score := range scores {
+			for top, ok := st.Peek(); ok && scores[top] < score; top, ok = st.Peek() {
+				st.Pop()
+
+				next[top] = i
+			}
+
+			if top, ok := st.Peek(); ok {
+				prev[i] = top
+			}
+
+			st.Push(i)
+		}
+
+		subarrays := make([]int64, len(nums))
+		for i := range subarrays {
+			subarrays[i] = int64(next[i]-i) * int64(i-prev[i])
+		}
+
+		return subarrays
+	}
+
+	scores := calculatePrimeScores()
+	subarrays := calculateSubarrays(scores)
+
+	q := pq.NewWith[int](subCmp)
+	for i := range nums {
+		q.Enqueue(i)
+	}
+
+	score := int64(1)
+	for k > 0 {
+		i, ok := q.Dequeue()
+		if !ok {
+			panic("queue should not become empty")
+		}
+
+		operations := min(int64(k), subarrays[i])
+
+		score = (score * power(int64(nums[i]), operations)) % MOD
+		k -= int(operations)
+	}
+
+	return int(score)
+}