hello-algo/codes/python/chapter_backtracking/subset_sum_i.py
2023-06-21 02:58:24 +08:00

48 lines
1.5 KiB
Python

"""
File: subset_sum_i.py
Created Time: 2023-06-17
Author: Krahets (krahets@163.com)
"""
def backtrack(
state: list[int], target: int, choices: list[int], start: int, res: list[list[int]]
):
"""回溯算法:子集和 I"""
# 子集和等于 target 时,记录解
if target == 0:
res.append(list(state))
return
# 遍历所有选择
# 剪枝二:从 start 开始遍历,避免生成重复子集
for i in range(start, len(choices)):
# 剪枝一:若子集和超过 target ,则直接结束循环
# 这是因为数组已排序,后边元素更大,子集和一定超过 target
if target - choices[i] < 0:
break
# 尝试:做出选择,更新 target, start
state.append(choices[i])
# 进行下一轮选择
backtrack(state, target - choices[i], choices, i, res)
# 回退:撤销选择,恢复到之前的状态
state.pop()
def subset_sum_i(nums: list[int], target: int) -> list[list[int]]:
"""求解子集和 I"""
state = [] # 状态(子集)
nums.sort() # 对 nums 进行排序
start = 0 # 遍历起始点
res = [] # 结果列表(子集列表)
backtrack(state, target, nums, start, res)
return res
"""Driver Code"""
if __name__ == "__main__":
nums = [3, 4, 5]
target = 9
res = subset_sum_i(nums, target)
print(f"输入数组 nums = {nums}, target = {target}")
print(f"所有和等于 {target} 的子集 res = {res}")