基于并行数组优先级的索引映射数组构建：线性时间无Python循环实现方案咨询

Absolutely! You can solve this problem in linear time without any Python-level loops by leaning into NumPy's vectorized operations and its ufunc.at functionality, which handles the update logic efficiently at the C level. Let's walk through this solution:

Problem Recap #

We have two arrays:

• a: An array of N integers (values range from 0 to M-1)
• b: An array of N positive numbers

We need to create array c (length M) where:

• c[i] is the index of the element in a equal to i that has the largest corresponding value in b
• If no such element exists, c[i] = -1

Example:
For N=5, M=3, a = [2, 1, 1, 2, 2], b = [1, 3, 5, 7, 3], the result is c = [-1, 2, 3]

Existing Approaches #

Let's quickly recap the two approaches you mentioned to frame the problem:

Approach 1: Python Loop Update

import numpy as np

c = -np.ones(M)
d = np.zeros(M)
for i, (idx, val) in enumerate(zip(a, b)):
    if d[idx] <= val:
        c[idx] = i
        d[idx] = val

• Time Complexity: O(N)
• Downside: Python loops are slow for large N, as they don't leverage NumPy's optimized C backend.

Approach 2: Sort-Based Vectorized Solution

import numpy as np

sort_idx = np.argsort(b)
a_idx = np.arange(len(a))
a_sorted = a[sort_idx]
a_idx_sorted = a_idx[sort_idx]
c = -np.ones(M)
c[a_sorted] = a_idx_sorted

• Upside: No Python loops
• Downside: Sorting adds O(N log N) time complexity, which is slower than linear time for large datasets.

Linear-Time, Loop-Free Solution #

Here's the optimized solution using NumPy's vectorized operations to achieve O(N) time without Python loops:

import numpy as np

def compute_target_c(a, b, M):
    # Initialize c with -1, and max_b to track the maximum b value for each a category
    c = -np.ones(M, dtype=int)
    max_b = np.full(M, -np.inf)  # Start with -inf since b values are positive
    
    # Use np.maximum.at to perform vectorized, in-place max updates (C-level operation)
    np.maximum.at(max_b, a, b)
    
    # Create a mask to find all indices where b matches the max_b for its a value
    matches_max = b == max_b[a]
    
    # Extract valid a values and their corresponding indices
    valid_a_values = a[matches_max]
    valid_indices = np.arange(len(a))[matches_max]
    
    # Assign indices to c: later indices overwrite earlier ones (matches Approach 1's logic)
    c[valid_a_values] = valid_indices
    
    return c

Test It with Your Example

a = np.array([2, 1, 1, 2, 2])
b = np.array([1, 3, 5, 7, 3])
M = 3
print(compute_target_c(a, b, M))  # Output: [-1  2  3]

How This Works #

1. np.maximum.at: This is the key here. It's a vectorized in-place operation that updates max_b at positions specified by a with the maximum of the current max_b value and the corresponding b value. This runs entirely in NumPy's optimized C code—no Python loop involved—so it's fast and linear time.
2. Masking for Valid Indices: We filter out all indices where the b value equals the maximum value for its category in a. This gives us all candidates that could be the answer for their respective a values.
3. Final Assignment: By assigning these valid indices to c using valid_a_values as the target positions, later indices overwrite earlier ones (just like in Approach 1). This ensures that if multiple indices have the same maximum b value, we keep the last one (which aligns with the problem's implicit logic when b values are tied).

Time Complexity #

Every step here runs in O(N) time:

• np.maximum.at processes each element once
• Masking and index extraction process each element once
• Final assignment processes the valid elements once

This gives us the best of both worlds: linear time and no slow Python loops.

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