如何优化HackerEarth Best Index问题的Java代码解决超时问题？

问题描述 #

给定一个包含N个元素的数组A，需选择最佳索引：若某个索引的特殊和为所有索引中的最大值，则该索引为最佳索引。
计算索引i的特殊和规则：

• 先取A[i]加入总和
• 接着取接下来2个元素A[i+1]、A[i+2]加入总和
• 再取接下来3个元素，以此类推，直到无法再取元素为止
  示例：数组有10个元素，选择索引3时，特殊和为A[3] + (A[4]+A[5]) + (A[6]+A[7]+A[8])（后续无法再取元素）。
  要求：找出最佳索引对应的最大特殊和，若有多个最佳索引，仅输出最大特殊和。

输入输出 #

• 输入：第一行是整数N，第二行是N个空格分隔的数组元素
• 输出：最大的特殊和

约束条件 #

• 1 ≤ N ≤ 10^5
• −10^7 ≤ A[i] ≤ 10^7

原代码（超时版本） #

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.*;
class TestClass {
    public static void main(String args[] ) throws Exception {        
        Scanner sc = new Scanner(System.in);
        ArrayList<Integer> list = new ArrayList<>();

        int n = sc.nextInt();
        int[] arr = new int[n];
        for(int i=0; i<n; i++){
            arr[i] = sc.nextInt();
        }
        list.add(0);
        for(int i=1;i<n; i++){
            int sum = list.get(i-1); ;
            list.add(sum+i);
        }
        int pos = 0;
        for(int i=0;i<list.size();i++){
            if(list.get(i)>n){
                pos = i-1;
                break;
            }
            else if(list.get(i)==n){
                pos = i;
                break;
            }
        }
        int len = n;
        int max = 0;
        for(int i=0; i< n;i++){
            int sum = 0;
            if(list.get(pos)>len){
                pos--;
            }
            for(int j=i;j<list.get(pos)+i;j++){
                sum = sum+arr[j];
            }
            len--;
            max = Math.max(max, sum);
        }
        System.out.println(max);
    }
}

超时原因分析 #

原代码时间复杂度为O(N²)：外层循环遍历每个索引（O(N)），内层循环逐个累加区间元素（最坏情况O(N)）。当N=10^{5时，总运算量达到10}10，远超时间限制。

优化思路与解决方案 #

核心优化是用前缀和数组快速计算区间和，将时间复杂度降至O(N√N)（完全满足10^5规模的运算要求）。

关键优化点 #

1. 前缀和数组：构建prefix数组，prefix[k]表示数组前k个元素的和（A[0]到A[k-1]）。区间A[l..r]的和可通过prefix[r+1] - prefix[l]直接计算，无需逐个累加。
2. 数学推导最大取数次数：对每个索引i，通过解方程找到最大的m，使得i + 1+2+...+m ≤ N（1+2+...+m = m*(m+1)/2），避免无效循环。
3. 输入优化：用BufferedReader替代Scanner，提升大规模输入的读取速度。
4. 类型安全：用long存储前缀和与特殊和，避免元素值过大导致的整数溢出。

优化后的代码 #

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class BestIndex {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int n = Integer.parseInt(br.readLine());
        long[] arr = new long[n];
        StringTokenizer st = new StringTokenizer(br.readLine());
        
        // 构建前缀和数组
        long[] prefix = new long[n + 1];
        prefix[0] = 0;
        for (int i = 0; i < n; i++) {
            arr[i] = Long.parseLong(st.nextToken());
            prefix[i + 1] = prefix[i] + arr[i];
        }
        
        long maxSum = Long.MIN_VALUE;
        
        for (int i = 0; i < n; i++) {
            // 计算当前索引i能取的最大m值
            long discriminant = 1 + 8L * (n - i);
            int m = (int)((Math.sqrt(discriminant) - 1) / 2);
            
            long currentSum = 0;
            int start = i;
            for (int k = 1; k <= m; k++) {
                int end = start + k;
                currentSum += prefix[end] - prefix[start];
                start = end;
            }
            
            maxSum = Math.max(maxSum, currentSum);
        }
        
        System.out.println(maxSum);
    }
}

复杂度说明 #

每个索引i对应的m值最大为√(2N)（约447当N=10^{5时），总运算量约4.47*10}7，完全符合时间限制。

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