多维度Top-k得分求解：探寻复杂度低于O(k^q)的通用算法

Hey, your current approach using partitions to narrow the search space is a solid start, but it's easy to see how O(k^q) complexity falls apart when q (number of columns) grows—think q=10 and k=10, that's 10 billion operations which is totally infeasible. Luckily, there are much more efficient strategies that bring complexity down to linear or low-polynomial levels relative to k and q. Let's break them down:

Method 1: Max-Heap Based Incremental Expansion (O(kqlogk)) #

This approach is similar to finding the top-K elements in a Cartesian product, using a max-heap (simulated with a min-heap by storing negative sums) to efficiently track and expand the best candidates:

1. Preprocess Columns: First, sort each column in descending order (so the largest elements are at the top) and keep track of their original indices in the input matrix.
2. Initialize Heap: Start with the combination where we take the largest element from every column (indices (0,0,...,0) in the sorted columns). Push this combination's sum (as a negative value, since Python's heapq is a min-heap) and its index tuple into the heap, plus mark it as visited to avoid duplicates.
3. Extract Top-K Elements:
   • Pop the largest sum from the heap (by popping the smallest negative value).
   • For each position in the current index tuple, generate a new candidate by incrementing that position's index (moving to the next-largest element in that column).
   • If the new candidate hasn't been visited and doesn't exceed the column length, calculate its sum, push it to the heap, and mark it as visited.
   • Repeat this k times to get your top-K combinations.

Code Implementation (Method 1) #

import numpy as np
import heapq

def get_top_k_combinations(matrix, k):
    N, q = matrix.shape
    # Sort each column descending, get sorted values and original indices
    sorted_cols = np.sort(matrix, axis=0)[::-1]
    original_col_indices = np.argsort(-matrix, axis=0)
    
    # Use min-heap to simulate max-heap (store negative sums)
    heap = []
    visited = set()
    
    # Start with the top element from each column
    initial_indices = tuple([0]*q)
    initial_sum = np.sum(sorted_cols[initial_indices, range(q)])
    heapq.heappush(heap, (-initial_sum, initial_indices))
    visited.add(initial_indices)
    
    top_k_results = []
    
    for _ in range(k):
        if not heap:
            break  # Exit early if there aren't enough combinations
        neg_sum, sorted_indices = heapq.heappop(heap)
        current_sum = -neg_sum
        # Map sorted indices back to original matrix indices
        original_indices = tuple(original_col_indices[idx, col] for col, idx in enumerate(sorted_indices))
        top_k_results.append((current_sum, original_indices))
        
        # Generate new candidates by incrementing each column's index
        for col_idx in range(q):
            new_sorted_indices = list(sorted_indices)
            new_sorted_indices[col_idx] += 1
            if new_sorted_indices[col_idx] >= N:
                continue  # Skip if we're beyond the column's length
            new_sorted_tuple = tuple(new_sorted_indices)
            if new_sorted_tuple not in visited:
                visited.add(new_sorted_tuple)
                new_sum = np.sum(sorted_cols[new_sorted_tuple, range(q)])
                heapq.heappush(heap, (-new_sum, new_sorted_tuple))
    
    return top_k_results

# Test with your example matrix
matrix = np.array([
    [0.9763, 0.0791, 0.1530, 0.5565],
    [0.1560, 0.1014, 0.6932, 0.7551],
    [0.8142, 0.9494, 0.4582, 0.4411],
    [0.3807, 0.2403, 0.6897, 0.7356],
    [0.0156, 0.9419, 0.9568, 0.2266]
])
k = 5
results = get_top_k_combinations(matrix, k)
for total, indices in results:
    print(f"{total:.4f} = q0[{indices[0]}] + q1[{indices[1]}] + q2[{indices[2]}] + q3[{indices[3]}]")

Running this code will output exactly the top-K sums you listed:

3.6376 = q0[0] + q1[2] + q2[4] + q3[1]
3.6301 = q0[0] + q1[4] + q2[4] + q3[1]
3.6181 = q0[0] + q1[2] + q2[4] + q3[3]
3.6106 = q0[0] + q1[4] + q2[4] + q3[3]
3.4755 = q0[2] + q1[2] + q2[4] + q3[1]

Method 2: Dynamic Programming with Heap Merging (O(qk²logk)) #

If you prefer a simpler implementation (especially for moderate q), this approach builds up the top-K sums incrementally by merging columns one by one:

1. Initialize: Start with the top-K elements from the first column (each element is a sum of just itself, plus its index).
2. Merge Columns Iteratively:
   • For each subsequent column, take the current top-K sums and combine them with the top-K elements of the new column (since combining with elements beyond the top-K can't make it into the new top-K).
   • Generate all k*k possible sums from these combinations, then filter down to the new top-K sums.
3. Final Result: After merging all columns, you'll have your top-K combination sums.

This is easier to code but slightly less efficient than Method 1, but still way better than O(k^q) for any reasonable k and q.

Complexity Comparison #

Extra Optimizations #

• For Method 1: Instead of using a full visited set, you can use a more memory-efficient structure like a hash set optimized for tuples, or even leverage the sorted nature of columns to avoid duplicates (though the set is simple and effective for most cases).
• For Method 2: When merging, use a min-heap to track the top-K sums on the fly instead of generating all k² sums first. This reduces memory usage from O(k²) to O(k) while keeping the same time complexity.

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