生成满足元素范围条件的r维所有可能向量的高效实现问询

Hey Gabriel, let's tackle this problem with a way more efficient approach—your current method works, but it's doing a ton of unnecessary work by generating far more vectors than you need and then filtering them out. Let's break down the better solutions:

The Optimal Itertools Approach #

Instead of generating all vectors up to the maximum value in d, you can directly create the valid range for each dimension and combine them with itertools.product. This way, you only produce vectors that meet your criteria from the start—no filtering required.

Here's the simplified, high-performance code:

import itertools

def generate_valid_vectors(d):
    # Create a range of valid values for each element (0 to d[i], inclusive)
    dimension_ranges = [range(0, num + 1) for num in d]
    # Combine all ranges to generate only valid vectors
    return [list(vec) for vec in itertools.product(*dimension_ranges)]

Why This Is Way Better #

• Efficiency: If your d has values like [100, 1], your original method would generate 101*101 = 10201 vectors just to filter down to 200 valid ones. This approach generates only the 200 valid vectors directly.
• Simplicity: No nested loops or manual validation counting—itertools.product handles all the combination logic for you.
• Memory Savings: You don't waste memory storing a massive list of invalid vectors before filtering.

Recursive Approach (If You Want to Avoid Itertools) #

Since you mentioned considering recursion, here's a clean recursive implementation that builds valid vectors step-by-step:

def generate_valid_vectors_recursive(d):
    # Base case: empty input returns a single empty vector
    if not d:
        return [[]]
    # Recursively generate valid vectors for the rest of the list
    rest_vectors = generate_valid_vectors_recursive(d[1:])
    # Combine each valid value from the first dimension with all rest vectors
    result = []
    for val in range(0, d[0] + 1):
        result.extend([[val] + vec for vec in rest_vectors])
    return result

How This Works #

• We split the problem into smaller subproblems: first handle the first element of d, then recursively solve for the remaining elements.
• For each valid value in the first dimension, we prepend it to every vector generated from the rest of the list, building up all valid combinations.

A Quick Tweak to Your Original Code #

If you ever need to adjust your original approach for small-scale cases, you can simplify the filtering step (no need for copy or manual counting):

import itertools

def main(d):
    greatest = max(d)
    all_poss_arr = itertools.product(range(greatest + 1), repeat=len(d))
    return [list(vec) for vec in all_poss_arr if all(v <= d[i] for i, v in enumerate(vec))]

But again, this is still inefficient compared to the first method because it generates unnecessary vectors.

Give the first method a go—it's the cleanest and most efficient way to solve this problem!

内容的提问来源于stack exchange，提问作者Gabriel Deza