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

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)
}