划分为k个相等的子集 -- 回溯算法应用

划分为k个相等的子集

class CanPartitionKSubsets:
    """
    698. 划分为k个相等的子集
    https://leetcode.cn/problems/partition-to-k-equal-sum-subsets/
    """
    def solution(self, nums: List[int], k: int) -> bool:
        """
        以桶的视角进行穷举，每个桶需要遍历nums中的所有数字，决定是否把当前数字装进桶中；
        装满一个桶后，装下一个桶
        :param nums:
        :param k:
        :return:
        """
        if k > len(nums):
            return False

        nums_sum = sum(nums)
        if nums_sum % k != 0:
            return False

        # 使用位图技巧，整数used的第i位的1/0表示used[i]的true/false
        self.used = 0
        self.target = nums_sum/k
        # 备忘录，存储used数组的状态
        self.memo = dict()
        self.nums = nums
        return self.backtrack(k, 0, 0)

    def backtrack(self, k, bucket, start):
        """

        :param k: k号桶
        :param bucket: k号桶已经装的数字之和
        :param start: 第start个元素
        :return:
        """
        if k == 0:
            # 所有桶都被装满了
            return True

        if self.used in self.memo.keys():
            return self.memo[self.used]

        if bucket == self.target:
            # 装满了当前桶，递归穷举下一个桶
            res = self.backtrack(k-1, 0, 0)
            self.memo[self.used] = res
            return res

        # 从 start 开始向后探查有效的 nums[i] 装⼊当前桶
        for i in range(start, len(self.nums)):
            # 判断第i位是否为1
            if (self.used >> i) & 1 == 1:
                continue

            if self.nums[i] + bucket > self.target:
                continue

            # 将第i位置为1
            self.used |= (1 << i)
            bucket += self.nums[i]

            #  递归穷举下⼀个数字是否装⼊当前桶
            if self.backtrack(k, bucket, i+1):
                return True

            # 将第i位置为0
            self.used ^= (1 << i)
            bucket -= self.nums[i]

        # 穷举了所有数字，都无法满足当前桶
        return False


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74