Python中寻找元素范围为0≤n≤end的6元素组合且和为指定值S的高效算法咨询

Hey there! Your problem boils down to finding ordered, bounded integer partitions—specifically, all 6-element lists where each element is between 0 and end (inclusive) and their sum equals S. Brute-force is way too slow (exponential complexity), so let's use a backtracking algorithm with smart pruning—this will drastically cut down unnecessary computations and get you the results efficiently.

Core Idea: Backtracking + Pruning #

Instead of checking every possible combination, we'll build the list one element at a time, and at each step we'll cut off paths that can't possibly lead to a valid solution. Here's the breakdown:
Suppose we've already picked k elements with a total sum of current_sum, and we have 6 - k elements left to choose:

1. The minimum possible sum of the remaining elements is 0 * (6 - k) (all zeros), and the maximum is end * (6 - k) (all end values).
2. Calculate how much more we need to reach S: remaining = S - current_sum.
3. If remaining is negative (we've already exceeded S) or remaining is larger than the maximum possible sum of the remaining elements (even filling the rest with end won't get us to S), we skip this entire branch—no need to waste time exploring it.

This pruning step is what makes this method way faster than brute force.

Pseudocode #

function generate_combinations(end, target_sum, current_position, current_list, current_total, result):
    if current_position == 6:
        if current_total == target_sum:
            add a copy of current_list to result
        return
    
    remaining_slots = 6 - current_position
    # Calculate the smallest number we can pick here: even if we fill the rest with end, we need at least this to hit target_sum
    min_num = max(0, target_sum - current_total - end * (remaining_slots - 1))
    # Calculate the largest number we can pick here: can't exceed end, and can't make current_total exceed target_sum
    max_num = min(end, target_sum - current_total)
    
    for num from min_num to max_num:
        add num to current_list
        generate_combinations(end, target_sum, current_position + 1, current_list, current_total + num, result)
        remove num from current_list

function f(end, S):
    result = empty list
    generate_combinations(end, S, 0, empty list, 0, result)
    return result

Python Implementation #

Here's a ready-to-use version that matches your function definition:

def f(end, S):
    result = []
    
    def backtrack(pos, current_list, current_sum):
        if pos == 6:
            if current_sum == S:
                result.append(current_list.copy())
            return
        
        remaining_slots = 6 - pos
        # Minimum number we can choose here to still have a shot at reaching S
        min_val = max(0, S - current_sum - end * (remaining_slots - 1))
        # Maximum number we can choose without exceeding S
        max_val = min(end, S - current_sum)
        
        for num in range(min_val, max_val + 1):
            current_list.append(num)
            backtrack(pos + 1, current_list, current_sum + num)
            current_list.pop()
    
    backtrack(0, [], 0)
    return result

# Test with your example
print(f(14, 14))

Why This Works So Well #

• The pruning logic eliminates entire branches of invalid possibilities before we even waste time exploring them. For example, if we're on the 3rd element and our current sum is already 15 (with S=14), we immediately stop that path—no need to check the remaining 3 elements.
• Unlike brute force's O(end^6) worst-case complexity, this method only traverses paths that could potentially lead to a valid solution. For your example f(14,14), it'll skip all combinations where the sum exceeds 14 early on, saving tons of time.

Bonus Tips for Scaling #

If you ever need to handle larger end values or longer lists:

• You can precompute the valid ranges for each position even further to narrow down the loop
• If you didn't care about ordered lists (but you do, per your example), you could generate unordered combinations first and then generate all permutations—but since your problem treats [14,0,...] and [0,14,...] as distinct, this isn't useful here.

内容的提问来源于stack exchange，提问作者Jingu Lee