LeetCode 🔴 410. Split Array Largest Sum

🔗 🔴 410. Split Array Largest Sum

tags: 動態規劃(DP) 記憶化搜索(Memoization) 前綴和(Prefix Sum) 二分搜尋(Binary Search)

題意：最小化最大值

給定一個非負整數陣列 $nums$ 和一個整數 $k$，你需要將這個陣列分成 $k$ 個非空的連續子陣列，使得 最大子陣列和 能夠被 最小化 。

返回分割後 最大子陣列和 的 最小值。

約束條件

• $1 \leq \text{nums.length} \leq 1000$
• $0 \leq \text{nums}[i] \leq 10^6$
• $1 \leq k \leq \min(50, \text{nums.length})$

思路

為求出每個子陣列的和，我們可以使用 前綴和(Prefix Sum) 來預先計算出陣列的部分和，這樣可以在 $O(1)$ 時間內獲得任意區間的和。令 $s[i]$ 表示前 $i$ 個數字的和，則 $s[i] = \text{nums}[0] + \text{nums}[1] + \ldots + \text{nums}[i-1] = \displaystyle\sum_{j=0}^{i-1} \text{nums}[j]$ ，則求 $[l, r]$ 區間的和，只需要計算 $s[r+1] - s[l]$ 即可。

根據題意，子陣列的最大值為 $\sum \text{nums}$ ，最小值為 $\max \text{nums}$ 。而根據約束條件， $\sum \text{nums} \leq 10^6 \times 1000 = 10^9$ ， $\max \text{nums} \leq 10^6$ 。

方法一：Top-down DP / 記憶化搜索(Memoization)

令 $f(i, j)$ 表示將前 $i$ 個數字分成 $j$ 個子陣列，所得到的最大和的最小值。則可以透過考慮上一段子陣列的結尾位置 $x$，將問題轉換為子問題 $f(x, j-1)$ ，此時的 最大子陣列和 為 $\max(f(x, j-1), s[i] - s[x])$ ，即子問題的 最大子陣列和 最小值、以及當前子陣列的總和，兩者的最大值。並可以將其作為答案的候選值，在所有可能的 $x$ 中，取最小值即可。

再來考慮一些細節和邊界情況：

• 對於轉移來源 $x$ 的範圍，由於子陣列必須為非空，因此前一段的結尾位置 $x$ 必須滿足 $0 \leq x < i$ 。
• 當 $j = 1$ 時，表示只有一個子陣列，此時可以直接返回前 $i$ 個數字的和 $s[i]$ ；
• 當 $j > i$ 時，由於每個子陣列需要至少包含一個數字，因此無法分成 $j$ 個子陣列，故此時返回一個足夠大的值，表示無法作為答案。

由於遞迴函數中存在大量重複計算的子問題，因此可以使用 記憶化搜索(Memoization) 來保存計算結果，避免重複計算。

最後返回 $f(n, k)$ 即可，其中 $n$ 為陣列的長度，表示考慮所有數字，分成 $k$ 個子陣列時，得到的 最大子陣列和 的最小值。

複雜度分析

• 時間複雜度：$O(n^2 \cdot k)$ ，其中 $n$ 為陣列 $nums$ 的長度，$k$ 為分割成的子陣列個數。
• 空間複雜度：$O(n \cdot k)$ ，記憶化搜索需要的空間。

PythonC++

若使用 functools.cache 進行記憶化搜索，可能需要一點運氣才能通過，可以多提交幾次試試看。

class Solution:
    def splitArray(self, nums: List[int], k: int) -> int:
        n = len(nums)
        s = list(accumulate(nums, initial=0))

        @cache
        def dfs(i: int, j: int) -> int: # 前 i 個數字分成 j 個子陣列，所得到最大和的最小值
            if j == 1:
                return s[i]
            if j > i: # 不能分成 j 個子陣列
                return 10**10
            res = 10**10
            for x in range(i): # 枚舉上一段的結尾是第 x 個數字
                res = min(res, max(dfs(x, j - 1), s[i] - s[x]))
            return res
        
        return dfs(n, k)

class Solution {
public:
    int splitArray(vector<int>& nums, int k) {
        int n = nums.size();
        vector<int> s(n + 1, 0); // prefix sum
        for (int i = 0; i < n; i++) s[i+1] = s[i] + nums[i];

        vector<vector<int>> memo(n + 1, vector<int>(k + 1, INT_MAX));
        function<int(int, int)> dfs = [&](int i, int j) -> int {
            if (j == 1) return s[i];
            if (j > i) return INT_MAX;
            int& res = memo[i][j];
            if (res != INT_MAX) return res;
            for (int x = 0; x < i; x++) {
                res = min(res, max(dfs(x, j - 1), s[i] - s[x]));
            }
            return res;
        };
        return dfs(n, k);
    }
};

方法二：Bottom-up DP

和方法一同樣的方式，只是將遞迴函數改為迭代，以 Bottom-up 的方式計算所有子問題的答案。類似地，我們可以使用一個二維陣列 $dp$ 來保存計算結果，其中 $dp[i][j]$ 表示將前 $i$ 個數字分成 $j$ 個子陣列，所得到的最大和的最小值。

• 初始化：$dp[0][0] = 0$ ，表示前 $0$ 個數字分成 $0$ 個子陣列，所得到的最大和的最小值為 $0$ ；對於其他位置，初始化為一個足夠大的值。
• 狀態轉移：對於每個 $i$ 和 $j$ ，可以枚舉上一段的結尾位置 $x$ ，並計算出子問題中 最大子陣列和 的最小值、以及當前子陣列的總和，兩者的最大值，在其中取最小值即可。因此，狀態轉移方程為 $dp[i][j] = \displaystyle\min_{x=0}^{i-1}(\max(dp[x][j-1], s[i] - s[x]))$ 。

使用三重循環來計算所有子問題的答案，最後返回 $dp[n][k]$ 即可。需要注意，在確定 $i$ 後，$j$ 的範圍需要在 $[1, \min(i, k)]$ 之間，否則無法將前 $i$ 個數字分成 $j$ 個子陣列。

複雜度分析

• 時間複雜度：$O(n^2 \cdot k)$ ，其中 $n$ 為陣列 $nums$ 的長度，$k$ 為分割成的子陣列個數。
• 空間複雜度：$O(n \cdot k)$ ，動態規劃表格 $dp$ 需要的空間。

PythonC++

class Solution:
    def splitArray(self, nums: List[int], k: int) -> int:
        n = len(nums)
        s = list(accumulate(nums, initial=0)) # prefix sum

        dp = [[10**10] * (k + 1) for _ in range(n + 1)]
        dp[0][0] = 0
        for i in range(1, n + 1):
            for j in range(1, min(i, k) + 1):
                for x in range(i):
                    dp[i][j] = min(dp[i][j], max(dp[x][j - 1], s[i] - s[x]))
        return dp[n][k]

class Solution {
public:
    int splitArray(vector<int>& nums, int k) {
        int n = nums.size();
        vector<int> s(n + 1, 0); // prefix sum
        for (int i = 0; i < n; i++) s[i+1] = s[i] + nums[i];

        vector<vector<int>> dp(n + 1, vector<int>(k + 1, INT_MAX));
        dp[0][0] = 0;
        for (int i = 1; i <= n; i++)
            for (int j = 1; j <= min(i, k); j++)
                for (int x = 0; x < i; x++)
                    dp[i][j] = min(dp[i][j], max(dp[x][j - 1], s[i] - s[x]));
        return dp[n][k];
    }
};

方法三：二分搜尋(Binary Search)

雖然我們可以通過動態規劃的方式來解決這個問題，但是我們也能將問題轉換為判斷問題，即「是否可以將陣列分成 $k$ 個子陣列，使得每個子陣列的和都小於等於 $d$」，如此便可以使用 二分搜尋(Binary Search) 來解決。

定義一個函數 check(d) 來檢查是否可以將陣列分成 $k$ 個子陣列，使得每個子陣列的和都小於等於 $d$。

• 利用 $cnt$ 來記錄 子陣列和 $\leq d$ 的子陣列個數，$s$ 來記錄當前子陣列的和。
• 遍歷陣列 $nums$ 中的每個元素 $x$ ，如果 $s + x > d$ ，則表示需要開始一個新的子陣列，並將 $s$ 設為 $x$ ；否則，將 $s$ 加上 $x$ 。
• 最後返回 $cnt \leq k$ ，即是否可以將陣列分成 $k$ 個子陣列，使得每個子陣列的和都小於等於 $d$ 。
  • 若可以分成 $k-1$ 個 子陣列和 都 $\leq d$ 的子陣列，則可以將其中某個子陣列再分成兩個子陣列，直到湊滿 $k$ 個子陣列為止，因此返回 $cnt \leq k$ 。
  • 之所以能夠這樣檢查，是因為題目確保了 $k \leq \text{nums.length}$ ，即 $k$ 不會超過陣列的長度，故一定可以分成 $k$ 個子陣列。

這裡使用左閉右閉區間進行二分搜尋，令 $left = \max \text{nums}$ ，$right = \sum \text{nums}$ ，在 $[left, right]$ 的範圍內進行二分搜尋，不斷調整左右範圍直到找到最大的最小價格差為止。

• 根據定義的 check 函數，區間的左側為 $\text{False}$ ，而右側為 $\text{True}$。
• 如果 $\text{check}(mid) = \text{True}$，則將右範圍縮小到 $[left, mid-1]$。
• 如果 $\text{check}(mid) = \text{False}$，則將左範圍擴大到 $[mid+1, right]$。
• 當區間為空，即 $left > right$ 時，終止搜尋。

最終返回 $left$ ，即 最小化 的 最大子陣列和。

複雜度分析

• 時間複雜度：$O(n \log s)$ ，其中 $s$ 表示陣列 $nums$ 中所有元素的和。
  • 每次二分搜尋時，需要對陣列進行一次遍歷，時間複雜度為 $O(n)$ ，因此總時間複雜度為 $O(n \log s)$ 。
• 空間複雜度：$O(1)$ 。

PythonC++

class Solution:
    def splitArray(self, nums: List[int], k: int) -> int:
        # 是否可以將 nums 分成 k 個子陣列，使得每個子陣列的和都小於等於 d
        def check(d: int) -> bool:
            s, cnt = 0, 1
            for x in nums:
                if s + x > d:
                    cnt += 1
                    s = x
                else:
                    s += x
            return cnt <= k
        
        left, right = max(nums), sum(nums)
        while left <= right:
            mid = (left + right) // 2
            if check(mid):
                right = mid - 1
            else:
                left = mid + 1
        return left

class Solution {
public:
    int splitArray(vector<int>& nums, int k) {
        
        auto check = [&](int d) -> bool {
            int s = 0, cnt = 1;
            for (int x : nums) {
                if (s + x > d) {
                    cnt++;
                    s = x;
                } else {
                    s += x;
                }
            }
            return cnt <= k;
        };

        int left = *max_element(nums.begin(), nums.end());
        int right = accumulate(nums.begin(), nums.end(), 0);
        while (left <= right) {
            int mid = (left + right) / 2;
            if (check(mid)) right = mid - 1;
            else left = mid + 1;
        }
        return left;
    }
};

---

寫在最後

Cover photo is generated by @たろたろ, thanks for their work!

先來猜明天的每日一題 🤔