ABC 123 D問題

• 計算量の工夫

昨日ABC123に参加してきました.

D問題

ボクはコンテスト中には解くことができず解説を見てACしました.
x , y , z 個数づつ存在しているa , b , c を全て足して全探索していては到底間に合いません.
そこで今回は k の値に制約がついていることに注目します.

k <= min ( 3000 , x * y * z )

k は大きくても3000にしかならないことがわかります.

3つの数字が出てきた時は2つを計算して, 残り一つを処理する. 

これが計算量を落とす最高の手段です. (エレガントな数学的解法を除くと) .
今回はまずはじめに a + b の値を計算して全通り出しておきます.この結果 a + bの配列をソートして昇順に並べた際に実際の答えとして使われる可能性があるのは, 大きい方から k 個までになります.
これを考えると abc の和として考えられるのは k * z 個となります. これは全部列挙するこてができ, ソートして昇順からk個を出力すれば答えとなります.

import java.util.Arrays;
import java.util.Scanner;

class Main {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int x = sc.nextInt();
        int y = sc.nextInt();
        int z = sc.nextInt();
        int k = sc.nextInt();
        long[] a = new long[x];
        long[] b = new long[y];
        long[] c = new long[z];
        for (int i = 0; i < x; i++)
            a[i] = sc.nextLong();
        for (int i = 0; i < y; i++)
            b[i] = sc.nextLong();
        for (int i = 0; i < z; i++)
            c[i] = sc.nextLong();
        Arrays.sort(a);
        Arrays.sort(b);
        Arrays.sort(c);
        long[] ab = new long[x * y];
        for (int i = 0; i < x; i++) {
            for (int j = 0; j < y; j++) {
                int pos = (i * y) + j;
                ab[pos] = a[i] + b[j];
            }
        }
        sc.close();
        Arrays.sort(ab);
        long[] abc = new long[3000 * z];
        for (int i = 0; i < Math.min(3000, x * y); i++) {
            for (int j = 0; j < z; j++) {
                int pos = (i * z) + j;
                abc[pos] = ab[(x * y) - 1 - i] + c[j];
            }
        }
        Arrays.sort(abc);
        for (int i = 0; i < k; i++) {
            System.out.println(abc[(3000 * z) - 1 - i]);
        }
    }
}