AGC 023-A Zero-Sum Ranges に失敗談を添えて(Kotlin/Java) - あるエンジニアのAtCoder奮闘記

irisruneです。令和と聞いてこの問題が出てくる人はすごいと思いました。今回は4提出分のコードがあるのでまた長いです。

約1年前の問題で、自分がAtCoderを始めて1月もしないうちに挑戦した~~そして玉砕した~~問題です。この記事で触れました累積和の存在を初めて認識した問題でもあります。

まずは、累積和を使うことだけ覚えていた今Kotlinで解いたコードを。

fun combination2(m : Long) : Long{
    return m * (m - 1L) / 2L
}

fun main(args: Array<String>){
    val n = readLine()!!.toInt()
    val la = readLine()!!.split(" ").map(String::toLong)

    val lCumSum = mutableListOf(0L)
    for(i in 1..n){
        lCumSum.add(lCumSum[i - 1] + la[i - 1])
    }
    lCumSum.sort()
    var retValue = 0L
    var elm = lCumSum[0]
    var cntSame = 1L
    for(i in 1..n){
        if(lCumSum[i] == elm){
            cntSame++
            if(i == n) retValue += combination2(cntSame)
        }
        else{
            retValue += combination2(cntSame)
            elm = lCumSum[i]
            cntSame = 1L
        }
    }
    println(retValue)
}

累積和を求めただけでは終わらない~~200点ではない~~問題なんですが、解説通りの手法に行き着く普通…のコードですね。というわけで、1年前の自分のコードを見てみたいと思います。

なお、当時はJavaを使っていました。まずは開始30分ほどで提出したこちらのコード。

import java.util.Scanner;

public class Main {

    public static void main(String[] args) {
        @SuppressWarnings("resource")
        Scanner sc = new Scanner(System.in);

        int n = Integer.parseInt(sc.next());

        int[] a = new int[n];
        for(int i = 0; i < n; i++) {
            a[i] = Integer.parseInt(sc.next());
        }

        int count = 0;
        for(int i = 0; i < n; i++) {
            int sum = 0;
            for(int j = i; j < n; j++) {
                sum += a[j];
                if(sum == 0) {
                    count++;
                }
            }
        }

        System.out.println(count);
    }
}

これは愚直に解いてますね。当然ですがこれでは $O(N^2)$ 時間かかってしまいます。

次に、コンテスト終了3分後に提出したコードを。

import java.util.Scanner;

public class Main {

    public static void main(String[] args) {
        @SuppressWarnings("resource")
        Scanner sc = new Scanner(System.in);


        int n = Integer.parseInt(sc.next());

        long[] a = new long[n];
        for(int i = 0; i < n; i++) {
            a[i] = Long.parseLong(sc.next());
        }

        int count = 0;
        long[] sumAbsPart = new long[n];
        long sumAbs = 0;
        for(int i = n - 1; i >= 0; i--) {
            sumAbsPart[i] = sumAbs;
            sumAbs += Math.abs(a[i]);
        }
        for(int i = 0; i < n; i++) {
            long sum = 0;
            for(int j = i; j < n; j++) {
                sum += a[j];
                if(sum == 0) {
                    count++;
                }else if(Math.abs(sum) > sumAbsPart[i]) {
                    break;
                }
            }
        }

        System.out.println(count);

    }
}

こちらは根本は変えずに、枝刈りをしてTLEを解消しようとしていますね。結果としてはむしろTLEが増えてしまいました。

最後に、おそらく解説を見た後で書いたであろうコードを。

import java.util.ArrayList;
import java.util.Collections;
import java.util.Scanner;

public class Main {

    public static void main(String[] args) {
        @SuppressWarnings("resource")
        Scanner sc = new Scanner(System.in);


        int n = Integer.parseInt(sc.next());

        long[] a = new long[n];
        for(int i = 0; i < n; i++) {
            a[i] = Long.parseLong(sc.next());
        }

        ArrayList<Long> s = new ArrayList<Long>();
        s.add((long) 0);
        for(int i = 0; i < n; i++) {
            s.add(s.get(i) + a[i]);
        }
        Collections.sort(s);

        int count = 0;
        for(int i = 0; i < n + 1; i++) {
            long base = s.get(i);
            for(int j = i + 1; j < n + 1; j++) {
                if(s.get(j) == base) {
                    count++;
                }else {
                    break;
                }
            }
        }
        System.out.println(count);

    }
}

解説を見て累積和を求める部分までは理解できたようですが、その後の処理がまずいですね。どういうわけか $_{M}C_2$ を求めるのに数式ではなく（二重目の）ループを用いています。結局これでは最悪計算量が $O(N^2)$ かかってしまいます（最初のコードに比べTLEは減りました）。