[LeetCode] 2818. Apply Operations to Maximize Score

https://leetcode.com/problems/apply-operations-to-maximize-score/description/

leetcode - Apply Operations to Maximize Score

문제 유형: 모노토닉 스택, 스택, 그리디, 정수론, 정렬

문제 난이도: Hard

 

문제

You are given an array nums of n positive integers and an integer k.

Initially, you start with a score of 1. You have to maximize your score by applying the following operation at most k times:

• Choose any non-empty subarray nums[l, ..., r] that you haven't chosen previously.
• Choose an element x of nums[l, ..., r] with the highest prime score. If multiple such elements exist, choose the one with the smallest index.
• Multiply your score by x.

Here, nums[l, ..., r] denotes the subarray of nums starting at index l and ending at the index r, both ends being inclusive.

The prime score of an integer x is equal to the number of distinct prime factors of x. For example, the prime score of 300 is 3 since 300 = 2 * 2 * 3 * 5 * 5.

Return the maximum possible score after applying at most k operations.

Since the answer may be large, return it modulo 10^9 + 7.

 

n길이의 양의정수배열과 정수 k를 받습니다.

처음에는 1점으로 시작합니다.

다음의 연산으로 점수를 최대화 해야합니다:

• 이전에 고른 적 없는 길이가 1 이상인 [l, ...., r]의 부분배열을 골라서 그 내에서 prime score가 가장 높은 요소를 골라 점수에 곱합니다.
• prime score는 서로 다른 소인수들의 개수입니다.

최대 k번의 연산으로 얻을 수 있는 최대 점수를 구하시오.

숫자가 너무 커질 수도 있으니 10^9+7로 나누시오.

 

풀이

영어로 보기:

https://leetcode.com/problems/apply-operations-to-maximize-score/solutions/6592477/precomputing-sieve-monotonic-stack-sorti-44bn/

점수를 최대화하기 위해서, 가장 큰 값을 가지는 요소를 포함하는 부분배열들을 선택하는 것이 중요하다.

 

• isPrime: isPrime[i]는 i가 소수인지에 대한 불리언 값이 저장
• primes: 100000이하의 소수가 담겨있음
• primescores: primescores[i]는 nums[i]의 서로 다른 소인수의 개수.
• left: nums[i]보다 prime score가 더 큰 가장 가까운 왼쪽 요소를 저장. 없다면 -1.
• right: nums[i]보다 prime score가 크거나 같은 가장 가까운 오른쪽 요소를 저장. 없다면 n.
• st: left와 right계산에 쓰이는 단조 스택(monotonic stack)
• candidates: {nums[i], i}를 내림차순으로, 같을경우 인덱스 오름차순으로 정렬한 리스트
• use: nums[i]가 주인공이 되는 부분배열의 개수
• ans: 최종 답안

 

1. 우선, prime score를 precompute해야 한다.
   • 에라토스테네스의 체를 만들면서 실제 소수들을 따로 배열에 넣는다.
   • 구한 소수들로 prime score를 구한다. (자세한 구현은 아래 코드 참조)
2. left[i]와 right[i]의 범위를 계산한다.
   • monotonic stack를 이용해서 각 요소가 해당 subarray에서 dominant해질 수 있는 좌우 최대 범위를 찾는 것이다
3. candidates에 {nums[i], i}를 담아서 nums[i]기준 내림차순, i기준 오름차순정렬한다.
4. 각 후보에 대해서, 해당 요소를 반드시 포함하는 부분배열을 만든다.
   • 좌범위 left[i], 우범위 right[i]이전까지 가능하므로
   • use = (idx - left[idx]) * (right[idx] - idx)가 되고, ans에 x^use를 누적하면 된다.
   • use가 클 수 있으므로, 거듭제곱을 분할 정복을 이용해서 빠르게 구해준다.
5. k가 0이 되면 멈추고 ans를 리턴한다.

 

코드

C++

#pragma GCC optimize("O3", "unroll-loops");
using ll = long long;
using pii = pair<int, int>;
const int MOD = 1e9 + 7;
vector<int> primes;
vector<bool> is_prime;
static const auto init = []() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    is_prime.resize(100001, true);
    is_prime[0] = false;
    is_prime[1] = false;
    for (int i = 2; i <= 100000; i++) {
        if (!is_prime[i]) continue;
        primes.emplace_back(i);
        for (int j = 2 * i; j <= 100000; j += i) {
            is_prime[j] = false;
        }
    }
    return 0;
}();

struct comp {
    bool operator()(const pii &a, const pii &b) {
        if (a.first == b.first) {
            return a.second < b.second;
        }
        return a.first > b.first;
    }
};

class Solution {
   public:
    ll powmod(ll base, ll exp) {
        ll res = 1;
        while (exp > 0) {
            if (exp % 2) res = (res * base) % MOD;
            base = (base * base) % MOD;
            exp /= 2;
        }
        return res;
    }

    int get_prime_score(int x) {
        int cnt = 0;
        for (int p : primes) {
            if (1LL * p * p > x) break;
            if (x % p == 0) {
                cnt++;
                while (x % p == 0) x /= p;
            }
        }
        if (x > 1) cnt++;
        return cnt;
    }

    int maximumScore(vector<int> &nums, int k) {
        int n = nums.size();
        vector<int> primescores(n, 0);
        ll ans = 1;

        for (int i = 0; i < n; i++) {
            primescores[i] = get_prime_score(nums[i]);
        }

        vector<int> left(n), right(n);
        stack<int> st;

        for (int i = 0; i < n; i++) {
            while (!st.empty() && primescores[st.top()] < primescores[i])
                st.pop();
            left[i] = st.empty() ? -1 : st.top();
            st.push(i);
        }

        while (!st.empty()) st.pop();

        for (int i = n - 1; i >= 0; i--) {
            while (!st.empty() && primescores[st.top()] <= primescores[i])
                st.pop();
            right[i] = st.empty() ? n : st.top();
            st.push(i);
        }

        vector<pii> candidates(n);
        for (int i = 0; i < n; i++) {
            candidates[i] = {nums[i], i};
        }
        sort(candidates.begin(), candidates.end(), comp());

        for (const auto &[num, idx] : candidates) {
            ll use = 1LL * (idx - left[idx]) * (right[idx] - idx);
            use = min(use, (ll)k);
            ans = (ans * powmod(num, use)) % MOD;
            k -= use;
            if (k == 0) break;
        }

        return ans % MOD;
    }
};

 

Go

const MOD int = 1e9 + 7

var (
	isInitialized bool
	isPrime       []bool
	primes        []int
)

func initSieve() {
	isPrime = make([]bool, 100001)
	for i := 2; i <= 100000; i++ {
		isPrime[i] = true
	}
	for i := 2; i <= 100000; i++ {
		if !isPrime[i] {
			continue
		}
		primes = append(primes, i)
		for j := i * 2; j <= 100000; j += i {
			isPrime[j] = false
		}
	}
}

func getPrimeScore(x int) int {
	res := 0
	for _, prime := range primes {
		if prime*prime > x {
			break
		}
		if x%prime == 0 {
			res++
			for x%prime == 0 {
				x /= prime
			}
		}
	}
	if x > 1 {
		res++
	}
	return res
}

func powmod(base, exp int) int {
	res := 1
	for exp > 0 {
		if exp%2 == 1 {
			res = (res * base) % MOD
		}
		base = (base * base) % MOD
		exp /= 2
	}

	return res
}

type Pair struct {
	first  int
	second int
}

type ArraySolver struct {
	n        int
	left     []int
	right    []int
	elements []Pair
}

func NewArraySolver(nums []int) *ArraySolver {
	n := len(nums)
	a := &ArraySolver{
		n:        n,
		left:     make([]int, n),
		right:    make([]int, n),
		elements: make([]Pair, n),
	}

	primescores := make([]int, n)
	for i, num := range nums {
		primescores[i] = getPrimeScore(num)
	}

	st := make([]int, 0)
	for i := 0; i < n; i++ {
		for len(st) > 0 && primescores[st[len(st)-1]] < primescores[i] {
			st = st[0 : len(st)-1]
		}
		if len(st) > 0 {
			a.left[i] = st[len(st)-1]
		} else {
			a.left[i] = -1
		}
		st = append(st, i)
	}

	st = st[:0]
	for i := n - 1; i >= 0; i-- {
		for len(st) > 0 && primescores[st[len(st)-1]] <= primescores[i] {
			st = st[0 : len(st)-1]
		}
		if len(st) > 0 {
			a.right[i] = st[len(st)-1]
		} else {
			a.right[i] = n
		}
		st = append(st, i)
	}

	for i, num := range nums {
		a.elements[i] = Pair{num, i}
	}
	sort.Slice(a.elements, func(i, j int) bool {
		if a.elements[i].first == a.elements[j].first {
			return a.elements[i].second < a.elements[j].second
		}
		return a.elements[i].first > a.elements[j].first
	})

	return a
}

func (a *ArraySolver) getScore(k int) int {
	res := 1

	for _, elem := range a.elements {
		num, idx := elem.first, elem.second
		l := a.left[idx]
		r := a.right[idx]
		use := min(k, (idx-l)*(r-idx))
		res = (res * powmod(num, use)) % MOD
		k -= use
		if k == 0 {
			break
		}
	}

	return res % MOD
}

func maximumScore(nums []int, k int) int {
	if !isInitialized {
		initSieve()
		isInitialized = true
	}
	a := NewArraySolver(nums)
	return a.getScore(k)
}