WUSTOJ 1349: TLE（Java）算法优化

题目链接：1349: TLE

Description

WH在刷题时，设计出了如下代码：

#include<stdio.h>
int main()
{
    int i, j, cnt, k, N, K, a[5555];
    scanf("%d%d", &N, &K);
    int ans = 0;
    for (i = 1; i <= N; i++)
    {
    	scanf("%d", &a[i]);					// 输入N个数
    }
    for (i = 1; i <= N; i++)
    {
    	for (j = i + 1; j <= N; j++)
    	{
			cnt = 0;
			for (k = i; k <= j; k++)		// 计算区间[i,j]的数的和
			{
				cnt += a[k];
			}
			if (cnt % K == 0)				// 和是K的整数倍
			{
				ans++;						// 结果加1
			}
    	}
    }
    printf("%d\n", ans);
}

上面的代码据说可能会TLE（不管你信不信，反正我是信了，不信你复制粘贴上去看看？）你的任务是重新设计一个更高效的程序帮助WH计算出ans的值

Input

第一行2个整数N K（N <= 5000，K > 0）

第二行N个数a[i]： 第i个数为a[i]（0 <= a[i] <= 10000）

Output

输出ans的值

Sample Input

5 1
1 2 3 4 5
5 2
5 4 2 1 3

Sample Output

10
5

分析💬

上述代码：计算N个数中，和是K的倍数的区间的个数。

数组sum下标从0开始，以第二组数据为例，sum如下：
第一次循环后：

j	0	1	2	3	4
sum[j]	5

ans = 0
第二次循环后：

j	0	1	2	3	4
sum[j]	9	4

ans = 0
第三次循环后：

j	0	1	2	3	4
sum[j]	11	6	2

ans = 1
第四次循环后：

j	0	1	2	3	4
sum[j]	12	7	3	1

ans = 2
第五次循环后：

j	0	1	2	3	4
sum[j]	15	10	6	4	3

ans = 5

代码💻

/**
 * Time 495ms
 * @author wowpH
 * @version 1.1
 * @date 2019年7月5日下午12:53:35
 * Environment:	Windows 10
 * IDE Version:	Eclipse 2019-3
 * JDK Version:	JDK1.8.0_112
 */

import java.util.Scanner;

public class Main {
	public static void main(String[] args) {
		int[] sum = new int[5000];					// 保存区间的和
		Scanner sc = new Scanner(System.in);
		while (sc.hasNext()) {
			int N = sc.nextInt();
			int K = sc.nextInt();
			int ans = 0;
			for (int i = 0; i < N; ++i) {
				sum[i] = sc.nextInt();
				for (int j = i - 1; j >= 0; --j) {	// 更新前面的和
					sum[j] = sum[j] + sum[i];
					if (0 == sum[j] % K) {			// 是K的整数倍
						ans = ans + 1;				// 结果加1
					}
				}
			}
			System.out.println(ans);
		}
		sc.close();
	}
}