【備忘録 AtCoder】アルゴリズムによって計算時間が大きく異なる二つのコード

問題

ABC281

実行が間に合わなかったコード

Main.java

import java.util.Scanner;

public class Main{

    public static void main(String[] args){

        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        long t = sc.nextLong();
        int[] time = new int[n];
        for(int i=0; i<n; i++){
            time[i] = sc.nextInt();
        }

        long sum = 0;
        int i = 0;
        //ここのループがボトルネック
        while(sum < t){
            i ++;
            sum += time[i%n];
        }
        
        int ans_num;
        if(i%n==0){
            ans_num = n;
        }else{
            ans_num = i%n;
        }


        System.out.println(ans_num + " " + (time[i%n]-(sum-t)));

    }

}

コードないのwhile文のループ数がとても多くなり，TLE（実行時間超過）が頻発してしまいました．

改善ver

Main.java

import java.util.Scanner;

public class Main{

    public static void main(String[] args){

        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        long t = sc.nextLong();
        int[] time = new int[n];
        for(int i=0; i<n; i++){
            time[i] = sc.nextInt();
        }

        long time_oneround = 0;
        for(int i=0; i<n; i++){
            time_oneround += time[i];
        }

        long time_rem = t % time_oneround;

        long sum = 0;
        int i = -1;
        while(sum < time_rem){
            i++;
            sum += time[i];
        }

        System.out.println((i+1) + " " + (time[i]-(sum-time_rem)));

    }

}

「プレイリスト全体を何周するかを最初に計算しておいて、残りを1曲ずつ進めればO(N)である」と考えてアルゴリズムを組めば，計算量を大きく減らすことができる．無事にAC．

参考ツイート

終わりに

アルゴリズムって大事だなぁと思いまいた（小並感）