Python递归求解数组数对问题遇大输入性能与递归深度问题

问题描述 #

给定包含n个整数的数组A = [a₁,a₂,...,aₙ]，需找出并输出所有满足i<j且aᵢ×aⱼ ≤ max(aᵢ,aᵢ₊₁,...,aⱼ)的(i,j)数对总数。

输入格式：

• 第一行输入数组长度n
• 第二行输入n个空格分隔的整数

约束条件：1≤n≤5×10⁵，1≤aᵢ≤10⁹

现有实现问题 #

采用递归方式实现的求解函数存在以下性能问题：

• 当数组含10万个元素时，因Python递归深度限制导致运行失败；
• 当数组含50万个元素且奇数索引元素为2时，出现时间超时问题。

现有代码 #

#!/bin/python3

import math
import os
import random
import re
import sys

#
# Complete the 'solve' function below.
#
# The function is expected to return a LONG_INTEGER.
# The function accepts INTEGER_ARRAY arr as parameter.
#

def solve(arr):
    # Write your code here
    if len(arr)<=1:
        return 0

    indx=arr.index(max(arr))

    x= solve(arr[0:indx+1] if indx != len(arr)-1 else arr[0:indx]) + solve(arr[indx+1:])

    arr=sorted(arr[0:indx+1])+sorted(arr[indx+1:])

    i=indx-1 if indx == len(arr)-1 else indx
    j=indx+1 if indx !=len(arr)-1 else indx

    while i>=0 and j<len(arr):
        p=arr[i]*arr[j]
        if p <= arr[indx]:
            x+=i+1
            j+=1
        else:
            i-=1

    return x

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')
    arr_count = int(input().strip())
    arr = list(map(int, input().rstrip().split()))
    result = solve(arr)
    fptr.write(str(result) + '\n')
    fptr.close()

优化方案与代码 #

核心优化点 #

1. 替换递归为迭代分治：避免Python递归深度限制，改用栈存储待处理的区间信息；
2. 复用排序结果：分治过程中保留左右区间的已排序数组，避免重复排序带来的额外开销；
3. 修正最大值查找逻辑：原代码仅返回第一个最大值的位置，改为遍历区间找任意最大值位置，避免多最大值场景下的错误；
4. 独立处理左右排序数组：跨最大值的数对计算直接使用左右已排序数组，无需合并，减少内存操作。

优化后的代码 #

#!/bin/python3

import os
import sys

def solve(arr):
    n = len(arr)
    if n <= 1:
        return 0
    
    total = 0
    # 栈元素：(区间起始索引, 区间结束索引)
    stack = [(0, n-1)]
    
    while stack:
        start, end = stack.pop()
        if start >= end:
            continue
        
        # 查找区间内最大值的位置
        max_val = arr[start]
        mid = start
        for i in range(start, end+1):
            if arr[i] > max_val:
                max_val = arr[i]
                mid = i
        
        # 处理左区间[start, mid-1]和右区间[mid+1, end]
        left_start, left_end = start, mid-1
        right_start, right_end = mid+1, end
        
        # 获取左右区间的排序数组
        sorted_left = sorted(arr[left_start:left_end+1]) if left_start <= left_end else []
        sorted_right = sorted(arr[right_start:right_end+1]) if right_start <= right_end else []
        
        # 双指针计算跨mid的有效数对：左区间元素*右区间元素 <= max_val
        i = len(sorted_left) - 1
        j = 0
        cnt = 0
        while i >= 0 and j < len(sorted_right):
            product = sorted_left[i] * sorted_right[j]
            if product <= max_val:
                cnt += i + 1
                j += 1
            else:
                i -= 1
        total += cnt
        
        # 将左右区间压入栈，保证先处理左区间
        if right_start <= right_end:
            stack.append((right_start, right_end))
        if left_start <= left_end:
            stack.append((left_start, left_end))
    
    return total

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')
    arr_count = int(input().strip())
    arr = list(map(int, input().rstrip().split()))
    result = solve(arr)
    fptr.write(str(result) + '\n')
    fptr.close()

极端场景进阶优化 #

针对5e5级别的超大数组，可预先用稀疏表预处理数组，将区间最大值查找时间从O(n)降至O(1)，预处理时间为O(n log n)，进一步压缩整体时间复杂度。

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