First, let's break down why your current code is way too slow: you're using a nested loop over two 100,000-element lists, which results in 10^10 operations—way beyond what Python can handle in under a second (Python typically processes ~1e6-1e7 operations per second). The fix here is to use mathematical simplification to eliminate the nested loop entirely.

Step 1: Analyze the Color Condition #

Your condition (i + n - 2) % 3 == k can be rewritten using modular arithmetic rules to make grouping easier:

• Red area: (i + n - 2) % 3 == 0 → (i + n) % 3 == 2
• White area: (i + n - 2) % 3 == 1 → (i + n) % 3 == 0
• Yellow area: (i + n - 2) % 3 == 2 → (i + n) % 3 == 1

Since (i + n) % 3 = (i%3 + n%3) %3, we can precompute sums of elements in w_list and h_list grouped by their index modulo 3, then combine those sums to get the total areas.

Step 2: Precompute Grouped Sums #

For both lists, calculate three sums based on the index's remainder when divided by 3:

• w0: Sum of w_list[i] where i %3 ==0
• w1: Sum of w_list[i] where i %3 ==1
• w2: Sum of w_list[i] where i %3 ==2
• Repeat this for h_list to get h0, h1, h2

Step 3: Calculate Areas Using Grouped Sums #

Use the modular combinations that satisfy each color's condition to compute totals in constant time:

• Red area = w0*h2 + w1*h1 + w2*h0 (all pairs where index sum mod3=2)
• White area = w0*h0 + w1*h2 + w2*h1 (all pairs where index sum mod3=0)
• Yellow area = w0*h1 + w1*h0 + w2*h2 (all pairs where index sum mod3=1)

Optimized Code #

def calculateAreas(w_list, h_list):
    # Calculate grouped sums for width list
    w0 = w1 = w2 = 0
    for idx, w in enumerate(w_list):
        mod = idx % 3
        if mod == 0:
            w0 += w
        elif mod == 1:
            w1 += w
        else:
            w2 += w
    
    # Calculate grouped sums for height list
    h0 = h1 = h2 = 0
    for idx, h in enumerate(h_list):
        mod = idx % 3
        if mod == 0:
            h0 += h
        elif mod == 1:
            h1 += h
        else:
            h2 += h
    
    # Compute each area using precomputed sums
    red_area = w0*h2 + w1*h1 + w2*h0
    white_area = w0*h0 + w1*h2 + w2*h1
    yellow_area = w0*h1 + w1*h0 + w2*h2
    
    # Return the tuple in the required order
    return (white_area, yellow_area, red_area)

Why This Works #

This approach cuts the time complexity from O(n²) to O(n)—for 100k elements, that's only 200k total operations. Even in Python, this will run in well under 1 second. We also removed the inefficient list.insert calls by directly constructing the tuple with the correct order.

Test It #

Running your test case:

longList = [i for i in range(100000)]
print(calculateAreas(longList, longList[::-1]))

This will output the exact same result as your original code, but in a fraction of the time.

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