698. Partition to K Equal Sum Subsets

Medium
Description

給定一個整數陣列 nums 和一個正整數 k，判斷是否能將陣列分成 k 個非空子集，使得每個子集的元素總和相等。

Example:

Input: nums = [4,3,2,3,5,2,1], k = 4 Output: true（[5], [1,4], [2,3], [2,3]）

Intuition#

核心思路：嘗試將每個數分配到 k 個桶中，每個桶目標為 sum/k，用回溯 + 剪枝搜索。

  • 總和必須能被 k 整除，否則無解
  • 每個桶的目標和 = total / k
  • 嘗試將每個數放入某個桶，超過目標則回溯
  • 關鍵優化：降序排序（大數先放，提早剪枝）、跳過相同桶值

Approaches#

1: 回溯 + 剪枝#

  • 概念: 嘗試將每個元素分配到 k 個桶，降序排序加速剪枝
  • 時間複雜度: O(k^n) 最差情況，但剪枝後遠小於此
  • 空間複雜度: O(n)
class Solution {
    fun canPartitionKSubsets(nums: IntArray, k: Int): Boolean {
        val total = nums.sum()
        if (total % k != 0) return false
        val target = total / k

        nums.sortDescending()
        if (nums[0] > target) return false

        val buckets = IntArray(k)

        fun backtrack(idx: Int): Boolean {
            if (idx == nums.size) {
                return buckets.all { it == target }
            }

            val seen = mutableSetOf<Int>()
            for (i in 0 until k) {
                if (buckets[i] + nums[idx] > target) continue
                if (buckets[i] in seen) continue
                seen.add(buckets[i])

                buckets[i] += nums[idx]
                if (backtrack(idx + 1)) return true
                buckets[i] -= nums[idx]
            }

            return false
        }

        return backtrack(0)
    }
}

⭐2: Bitmask DP#

  • 概念: 用 bitmask 表示已選元素的狀態，DP 記錄每個狀態下當前桶的累積和，判斷是否能填滿所有桶
  • 時間複雜度: O(n * 2^n)
  • 空間複雜度: O(2^n)
class Solution {
    fun canPartitionKSubsets(nums: IntArray, k: Int): Boolean {
        val total = nums.sum()
        if (total % k != 0) return false
        val target = total / k

        val n = nums.size
        val dp = IntArray(1 shl n) { -1 }
        dp[0] = 0 // 空狀態，當前桶累積 = 0

        for (mask in 0 until (1 shl n)) {
            if (dp[mask] == -1) continue

            for (i in 0 until n) {
                if (mask and (1 shl i) != 0) continue // 已選
                val newMask = mask or (1 shl i)
                if (dp[newMask] != -1) continue // 已訪問

                val currentBucket = dp[mask] % target
                if (currentBucket + nums[i] <= target) {
                    dp[newMask] = dp[mask] + nums[i]
                }
            }
        }

        return dp[(1 shl n) - 1] == total
    }
}

🔑 Takeaways#

  • Pattern: K 等分問題 → 回溯分桶 or Bitmask DP
  • 關鍵技巧: 回溯法中降序排序 + 跳過相同桶值是效能關鍵；Bitmask DP 用 dp[mask] % target 判斷當前桶的填充量，滿了自動進入下一個桶。473. Matchsticks to Square 是此題 k=4 的特例