D - Cake 123

はじめに

競技プログラミングAtCoder Beginner Contest 123
D - Cake 123
言語はC#
何かあればTwitter→@pirorirori_n712まで

問題

問題文

AtCoder 洋菓子店は数字の形をしたキャンドルがついたケーキを販売しています。
ここには 1,2,3 の形をしたキャンドルがついたケーキがそれぞれ X 種類、Y 種類、Z 種類あります。
それぞれのケーキには「美味しさ」という整数の値が以下のように割り当てられています。

・1 の形のキャンドルがついたケーキの美味しさはそれぞれ A1,A2,...,AX
・2 の形のキャンドルがついたケーキの美味しさはそれぞれ B1,B2,...,BY
・3 の形のキャンドルがついたケーキの美味しさはそれぞれ C1,C2,...,CZ
高橋君は ABC 123 を記念するために、1,2,3 の形のキャンドルがついたケーキを 1 つずつ買うことにしました。
そのようにケーキを買う方法は X×Y×Z 通りあります。

これらの選び方を 3 つのケーキの美味しさの合計が大きい順に並べたとき、
1,2,...,K 番目の選び方でのケーキの美味しさの合計をそれぞれ出力してください。

制約

・1≤X≤1000
・1≤Y≤1000
・1≤Z≤1000
・1≤K≤min(3000,X×Y×Z)
・1≤Ai≤10,000,000,000
・1≤Bi≤10,000,000,000
・1≤Ci≤10,000,000,000
・入力中の値はすべて整数である。

入力

入力は以下の形式で標準入力から与えられる。

X Y Z K
A1 A2 A3 ... AX
B1 B2 B3 ... BY
C1 C2 C3 ... CZ

出力

i 行目に、問題文中の i 番目の値を出力せよ。

解き方

以下を参考にしています。
https://img.atcoder.jp/abc123/editorial.pdf
提出 #4865465 - AtCoder Beginner Contest 123

解き方の流れ

以下では、解説の中の「解法 #1 – 工夫したソートで 𝑶(𝑲^𝟐𝐥𝐨𝐠 𝑲)」で解いています。

ケーキの選び方のパターンを全探索し
おいしさの大きい順番に並び替えれば、この問題は解けますが
時間制限に引っかかってしまいます。

そのため、少し計算量を減らす工夫が必要になります。
全探索をする時間はないので、工程を以下のように分けます。

① A+Bのケーキの組み合わせのパターンを全て出す
② ①で出した組み合わせのパターンを「美味しさ」が大きい順に並び替える
③ ①で出したA+Bの組み合わせのパターンに、Cのケーキの組み合わせを足す


① A+Bのケーキの組み合わせのパターンを全て出す
A+Bのケーキの組み合わせのパターンを全て出し、配列に保存します。

入力の変数の保存

var line = Console.ReadLine().Split(' ').Select(x => Int32.Parse(x)).ToArray();
var X = line[0];
var Y = line[1];
var Z = line[2];
var K = line[3];
var A = Console.ReadLine().Split(' ').Select(x => Int64.Parse(x)).ToArray();
var B = Console.ReadLine().Split(' ').Select(x => Int64.Parse(x)).ToArray();
var C = Console.ReadLine().Split(' ').Select(x => Int64.Parse(x)).OrderByDescending(x => x).ToArray();
        

A+Bのケーキの組み合わせのパターンを全て出し、配列に保存

//A+Bのケーキのおいしさを保存する配列
var array = new long[X * Y];
var index = 0;
//Aのケーキ
for (int i = 0; i < X; ++i)
{
    //Bのケーキ
    for (int j = 0; j < Y; ++j)
    {
        array[index] = A[i] + B[j];
        index++;
    }
}


② ①で出した組み合わせのパターンを「美味しさ」が大きい順に並び替える
①で配列に保存したケーキの組み合わせのパターンを「美味しさ」が大きい順に並び替えます。

//A+Bのケーキの組み合わせのおいしさを保存する配列を
//おいしさが大きい順に並び替える
Array.Sort(array);
Array.Reverse(array);

この問題で「OrderByDescending().ToArray()」を使用すると
TLEしてしまうので、
「Array.Sort(),Array,Reverse()」で並び替えています。

③ ①で出したA+Bの組み合わせのパターンに、Cのケーキの組み合わせを足す

//A+BのおいしさにCのおいしさを足すループの長さ
//なるべくループしないように、KかX*Yの小さいほうをとる
var arrayLength = Math.Min(K, X * Y);
//A+B+Cのケーキのおいしさ
var arrayAll = new long[arrayLength * Z];
index = 0;
//A+BにCのおいしさを足す
for (int i = 0; i < arrayLength; ++i)
{
    //Cのケーキ
    for (int j = 0; j < Z; ++j)
    {
        arrayAll[index] = array[i] + C[j];
        index++;
    }
}

Cのケーキの組み合わせを足すループをなるべく少なくするために
KかX*Yの小さいほうをループする回数にしています。

全探索するにはX×Yになりますが
この問題では「1~K番目の美味しさ」を出力すればよく、
A+Bの「美味しさ」にCの「美味しさ」を足した際の全体の「美味しさ」で
「1~K番目に入る美味しさ」は
A+Bの「美味しさ」では必ずK番目以内におさまるので
ループの回数を抑えることができます。

出力例2で説明すると、

出力例2

3 3 3 5
1 10 100
2 20 200
1 10 100

このときA+Bの配列は以下のようになります。

array={300,210,201,120,102,30,21,12,3}

Kは5なので
大きいほうから数えて5番目の102までが
Cを足すときの対象になります。

もしここで、
5個よりも小さい個数、例えば2個を
Cを足すときの対象とした場合
Cのケーキが一つしかなかったら
美味しさを1~2までしか出すことができません。

逆に、
5個よりも大きい個数、例えば8個を
Cを足すときの対象とした場合
Cのケーキが一つだとしても
美味しさを1~8まで出すことになってしまい
3個余分になってしまいます。

そのため、
なるべく無駄なく計算するために
ループの回る回数を
KとX×Yの小さい方としています。

最後に、
A+B+Cのケーキの美味しさを格納している配列を
美味しさが大きい順に並び替え、
K個出力します。

//A+B+Cのケーキの組み合わせのおいしさを保存する配列を
//おいしさが大きい順に並び替える
Array.Sort(arrayAll);
Array.Reverse(arrayAll);
//K個出力
index = 0;
foreach (var ans in arrayAll)
{
    Console.WriteLine(ans);
    if (K <= ++index) break;
}

全体のコード

using System;
using System.Linq;
using System.Collections.Generic;

class Program
{
    static void Main()
    {
        var line = Console.ReadLine().Split(' ').Select(x => Int32.Parse(x)).ToArray();
        var X = line[0];
        var Y = line[1];
        var Z = line[2];
        var K = line[3];
        var A = Console.ReadLine().Split(' ').Select(x => Int64.Parse(x)).ToArray();
        var B = Console.ReadLine().Split(' ').Select(x => Int64.Parse(x)).ToArray();
        var C = Console.ReadLine().Split(' ').Select(x => Int64.Parse(x)).OrderByDescending(x => x).ToArray();
        //A+Bのケーキのおいしさを保存する配列
        var array = new long[X * Y];
        var index = 0;
        //Aのケーキ
        for (int i = 0; i < X; ++i)
        {
            //Bのケーキ
            for (int j = 0; j < Y; ++j)
            {
                array[index] = A[i] + B[j];
                index++;
            }
        }
        //A+Bのケーキの組み合わせのおいしさを保存する配列を
        //おいしさが大きい順に並び替える
        Array.Sort(array);
        Array.Reverse(array);
        //Cのケーキの組み合わせを足すループの長さ
        //なるべくループしないように、KかX*Yの小さいほうをとる
        var arrayLength = Math.Min(K, X * Y);
        //A+B+Cのケーキのおいしさ
        var arrayAll = new long[arrayLength * Z];
        index = 0;
        for (int i = 0; i < arrayLength; ++i)
        {
            //Cのケーキ
            for (int j = 0; j < Z; ++j)
            {
                arrayAll[index] = array[i] + C[j];
                index++;
            }
        }
        //A+B+Cのケーキの組み合わせのおいしさを保存する配列を
        //おいしさが大きい順に並び替える
        Array.Sort(arrayAll);
        Array.Reverse(arrayAll);
        //K個出力
        index = 0;
        foreach (var ans in arrayAll)
        {
            Console.WriteLine(ans);
            if (K <= ++index) break;
        }
    }
}