Wunder Fund 2016-F : Double Knapsack

F.java

import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.InputMismatchException;

public class F {
    public static void main(String[] args) {
        InputReader in = new InputReader(System.in);
        PrintWriter out = new PrintWriter(System.out);

        int n = in.nextInt();
        long[] a = new long[n];
        long[] b = new long[n];
        for (int i = 0; i < n ; i++) {
            a[i] = in.nextInt();
        }
        for (int i = 0; i < n ; i++) {
            b[i] = in.nextInt();
        }
        long sumA = 0;
        long sumB = 0;
        for (int i = 0; i < n ; i++) {
            sumA += a[i];
            sumB += b[i];
        }
        int[] segments = sumA <= sumB ? solve(a, b) : solve(b, a);
        int fromA = segments[0];
        int toA = segments[1];
        int fromB = segments[2];
        int toB = segments[3];
        if (sumA > sumB) {
            fromA = segments[2];
            toA = segments[3];
            fromB = segments[0];
            toB = segments[1];
        }

        String lineA = buildIntegers(fromA, toA);
        String lineB = buildIntegers(fromB, toB);

        out.println(toA - fromA + 1);
        out.println(lineA);
        out.println(toB - fromB + 1);
        out.println(lineB);
        out.flush();
    }

    private static int[] solve(long[] a, long[] b) {
        // Σa <= Σb
        int[] diffA = new int[1000010];
        int[] diffB = new int[1000010];
        Arrays.fill(diffA, -2);
        Arrays.fill(diffB, -2);
        diffA[0] = -1;
        diffB[0] = -1;

        int n = a.length;
        int fromA = -1;
        int toA = -1;
        int fromB = -1;
        int toB = -1;
        int bi = 0;
        long asum = 0, bsum = 0;
        for (int i = 0 ; i < n ; i++) {
            asum += a[i];
          
            // asumを超えないように足せるだけ足す
            while (bi < n && bsum + b[bi] <= asum) {
                bsum += b[bi];
                bi++;
            }
          
            // 差分はd (0 <= d <= n-1)
            int d = (int)(asum - bsum);

            // すでに差がdになるペアを見つけてたら終了。鳩の巣原理よりこのif文はどこかで真になる
            if (diffA[d] != -2) {
                fromA = diffA[d]+1;
                toA = i;
                fromB = diffB[d]+1;
                toB = bi-1;
                break;
            }
          
          　　 // 差がdになるペアを記憶
            diffA[d] = i;
            diffB[d] = bi-1;
        }
        return new int[]{fromA, toA, fromB, toB};

    }

    private static String buildIntegers(int fromA, int toA) {
        StringBuilder line = new StringBuilder();
        for (int i = fromA ; i <= toA ; i++) {
            line.append(' ').append(i+1);
        }
        return line.substring(1);
    }

    // 以下テンプレにつき略
}