[LeetCode] 3578. Count Partitions With Max-Min Difference at Most K

https://leetcode.com/problems/count-partitions-with-max-min-difference-at-most-k/description/?envType=daily-question&envId=2025-12-06

LeetCode - Count Partitions With Max-Min Difference at Most K

문제 유형: 모노토닉 큐, 다이나믹 프로그래밍, 슬라이딩 윈도우

문제 난이도: Medium

 

문제

You are given an integer array nums and an integer k. Your task is to partition nums into one or more non-empty contiguous segments such that in each segment, the difference between its maximum and minimum elements is at most k.

Return the total number of ways to partition nums under this condition.

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

 

정수 배열 nums와 정수 k가 주어집니다.

당신이 할 일은 1개 이상의 부분배열들로 나눠서 그 안의 최대 및 최소 요소의 크기가 최대 k이도록 해야합니다.

이렇게 나누는 모든 경우의 수를 구하시오.

정답이 클 수 있으니, 1e9+7로 나누시오.

 

풀이

dp[i] = [0, n-1]구간에 대해서 구한 정답(1 <= i <= n). 그래서 기저 상태로 dp[0] = 1로 한다.

인덱스 i를 현재 만드는 구간의 마지막 부분이라고 생각하자.

그러면, 맨 앞이 j라고 치면, [j, ..., i]이다. 이 구간에서 [j, ..., i]의 max - min <= k를 유지해야 한다.

그리고, [0, ..., j-1]가 잘 나눠있어야 하고, 그 경우의 수는 dp[j]개이다.

 

큰 수들을 담는 모노토닉 데크 maxQ, 작은 수들을 담는 모노토닉 데크 minQ를 만든다.

i를 오른쪽 포인터 끝으로 두고, i를 추가하며 maxQ, minQ를 업데이트한 뒤, 왼쪽 끝인 j를 줄여보는식이다.

즉, 아래의 i in [0 ... n-1]동안 다음을 하면 된다:

- nums[maxQ의 오른쪽 끝]이 이번 nums[i]보다 작거나 같은 동안, maxQ의 오른쪽을 pop한다.

- nums[minQ의 왼쪽 끝]이 이번 nums[i]보다 크거나 같은 동안, minQ의 오른쪽을 pop한다.

- 이제, 최대 값이나 최소 값을 덜어내야 한다. 만약 최대 값의 차나 최소 값의 차가 k 초과라면, maxQ의 왼쪽 또는 minQ의 왼쪽 중, j를 가진 곳을 한칸 당겨서 계속 덜어내기를 시도한다.

이제, dp[i+1] = dp[j] + dp[j+1] + ... + dp[i]로 구해주면 되는데. 

이를 더 빨리 구하기 위해 prefix sum 배열을 유지해주면 된다.

 

마지막으로, dp[n]을 반환하면 된다.

 

 

코드

C++

#pragma GCC optimize("O3", "unroll-loops");
static const int __ = [](){
    ios::sync_with_stdio(0);
    cin.tie(0);
    return 0;
}();
using ll = long long;
const ll MOD = 1e9+7;
class Solution {
public:
    int countPartitions(vector<int>& nums, int k) {
        int n = nums.size();
        vector<ll> dp(n+1);
        vector<ll> pref(n+1);
        deque<int> minQ, maxQ;

        dp[0] = 1;
        pref[0] = 1;
        for(int i = 0, j = 0; i < n; i++) {
            while(!maxQ.empty() && nums[maxQ.back()] <= nums[i]) {
                maxQ.pop_back();
            }
            maxQ.push_back(i);
            while(!minQ.empty() && nums[minQ.back()] >= nums[i]) {
                minQ.pop_back();
            }
            minQ.push_back(i);

            while(!maxQ.empty() && !minQ.empty() && nums[maxQ.front()] - nums[minQ.front()] > k) {
                if(maxQ.front() == j) {
                    maxQ.pop_front();
                }
                if(minQ.front() == j) {
                    minQ.pop_front();
                }
                j++;
            }
            dp[i+1] = (pref[i] - (j > 0 ? pref[j - 1] : 0) + MOD) % MOD;
            pref[i+1] = (pref[i] + dp[i+1]) % MOD;
        }

        return dp[n];
    }
};