AtCoder Beginner Contest 123 C,D問題メモ
C - Five Transportations
問題
- 都市1~6を、1から6まで順に移動する
- 移動は乗り物に乗って行い、交通手段は以下の通り
- 都市1→2の移動は定員 $A$ 人
- 都市2→3の移動は定員 $B$ 人
- 都市3→4の移動は定員 $C$ 人
- 都市4→5の移動は定員 $D$ 人
- 都市5→6の移動は定員 $E$ 人
- 各都市間の移動は全て1分かかる
- 各交通手段は、全て時刻0から1分間隔で出発する
- 時刻0に $N$ 人が都市1にいる
- 全員が都市6に到着する最短時間を求めよ
解法
一番定員の少ない交通手段がボトルネックとなる。
たとえば都市3→4の移動(定員 $C$ 人)がボトルネックだったとする。
すると、$A,B,D,E$ の値にかかわらず、全ての都市間で $C$ 人ずつ動いたとしても、全体の到着時刻は変わらない。
N=10 都市 1 2 3 4 5 6 定員 5 4 3 6 7 [動けるだけ動いた場合] t=0 10 t=1 5 5 t=2 6 4 t=3 2 5 3 t=4 4 3 3 t=5 1 3 3 3 t=6 1 3 6 t=7 1 9 t=8 10 [全て3人ずつ動いた場合] t=0 10 t=1 7 3 t=2 4 3 3 t=3 1 3 3 3 t=4 1 3 3 3 t=5 1 3 3 3 t=6 1 3 6 t=7 1 9 t=8 10
都市3→4間を $N$ 人が通過するのに $\displaystyle ceil(\frac{N}{C})$ 分かかる。
それ以外に、先頭が都市1から3まで到着するのに2分、末尾が都市4から6まで移動するのに2分かかる。 これは、ボトルネックがどこであろうと、合計4分かかる。(5つの都市間の内、ボトルネック以外の4箇所の移動時間)
よって答えは、$\displaystyle ceil(\frac{N}{C})+4$ 分。
import math n = int(input()) times = [int(input()) for _ in range(5)] print(int(math.ceil(n / min(times))) + 4)
D - Cake 123
問題
- あるケーキ店では、ケーキに3種類のキャンドルのいずれか1つをつけて販売している
- キャンドル1,2,3のついたケーキがそれぞれ $X,Y,Z$ 個ある
- それぞれのケーキの美味しさが整数値でわかっており、
- キャンドル1のケーキは $A_1,A_2,...,A_X$
- キャンドル2のケーキは $B_1,B_2,...,B_Y$
- キャンドル3のケーキは $C_1,C_2,...,C_Z$
- キャンドル1,2,3のついたケーキからそれぞれ1つずつ、計3つを選ぶことを考える
- そのような選び方は $X \times Y \times Z$ 通りある
- ある選び方の得点を、3つのケーキの美味しさの合計とする
- 得点を大きい方から並べたとき、$1~K$ 番目までの、それぞれの得点を求めよ
- $1 \le X,Y,Z \le 1000$
- $1 \le K \le \min(3000, X \times Y \times Z)$
解法
$K$ の制約の書き方がちょっと怪しい。最大でも3000までである。
まず、愚直に全ての選び方の得点を求めていると、$XYZ=10^9$ オーダーの計算が必要となり、間に合わない。
普通に考えると得点の低い方はわざわざ計算しなくていいが、それをどうやって限定するかがポイントとなる。
以下、キャンドル1,2,3の刺さったケーキをそれぞれ「$A$ のケーキ」「$B$ のケーキ」「$C$ のケーキ」と呼ぶ。
2組での全探索を2回やる
3組を同時に考えるから遅いんであって、$A$ と $B$ の2個の組み合わせを全て求めるだけなら、$XY=10^6$ オーダーとなり、間に合う範疇である。
$A_i+B_j$ の中で上位 $K$ 位以内に入れなかった $(i,j)$ は、 最終的な $A_i+B_j+C_k$ の順位でも上位 $K$ 個に入る可能性は無い。
なので、$X \times Y$ 個の候補から上位 $K$ 個に絞り込んでよくなる。
絞り込んだ上で改めてそれと $C$ の組み合わせを求める。オーダーは $KZ = 3 \times 10^6$ となり、こちらも間に合う。
高速化
Pythonは、純粋に $A$ と $B$ から $XY$ 要素のリストを作りソートしていると、制限時間ぎりぎりか、オーバーしてしまう。
heapq.nlargest()
イテレータから大きい順に上位 $K$ 個の要素を取得するのに打って付けの heapq.nlargest()
がある。
基本的にこういう組み込みライブラリはC言語で実装されているので、利用できるものは利用すると高速になる。
事前ソート
または、あらかじめ $A,B,C$ を個別にソートしておく。
Pythonのsortの内部実装はTimSortと呼ばれ、マージソートみたいな感じらしい。
これは、データ全体がある程度ソート済みだとスキップが多く発生して理論上の平均計算量より少なくなる性質がある。
# PythonだとTLEかも x, y, z, k = map(int, input().split()) A = list(map(int, input().split())) B = list(map(int, input().split())) C = list(map(int, input().split())) AB = [a + b for a in A for b in B] AB.sort(reverse=True) AB = AB[:k] ABC = [ab + c for ab in AB for c in C] ABC.sort(reverse=True) ABC = ABC[:k] print('\n'.join(map(str, ABC)))
from heapq import nlargest x, y, z, k = map(int, input().split()) A, B, C = [list(map(int, input().split())) for _ in [0] * 3] AB = nlargest(k, (a + b for a in A for b in B)) ABC = nlargest(k, (ab + c for ab in AB for c in C)) print('\n'.join(map(str, ABC)))
x, y, z, k = map(int, input().split()) A, B, C = [sorted(map(int, input().split()), reverse=True) for _ in [0] * 3] AB = sorted((a + b for a in A for b in B), reverse=True)[:k] ABC = sorted((ab + c for ab in AB for c in C), reverse=True)[:k] print('\n'.join(map(str, ABC)))
優先度キューを用いる
$A,B,C$ の各ケーキを美味しい順にソート済みとし、 $A_i,B_i,C_i$ を、それぞれのキャンドルで $i$ 番目に美味しいケーキの美味しさとする。
ある組み合わせ $(A_i,B_j,C_k)$ が上位 $K$ 番以内に入っているとする。
すると、この組合せから $A,B,C$ のいずれか1つを変えた組合せで次に得点が高くなるのは、以下のどれかである。
$$(A_{i+1},B_j,C_k),(A_i,B_{j+1},C_k),(A_i,B_j,C_{k+1})$$
優先度キューから最も得点の大きいものを取りだし、次に得点が高い可能性のある最大3つを優先度キューに加える、というのを $K$ 回繰り返せばよい。
1x1x1の立方体が積み重なった $X \times Y \times Z$ の直方体において、(1,1,1)から、面で隣接する3方向に徐々に探索を広げていく様子をイメージするとよいかもしれない。
簡単のために $A,B$ の2次元で考える。盤面は優先キューの中身のイメージを表す。
■: K番以内確定(取りだし済み) ○: 優先キューに入っている、数字は得点 それ以外: 未探索 [初期状態] 各ケーキの最大値の合計は明らかに1番得点が高いのでここからスタート K = 5 A1 A2 A3 A4 10 7 3 2 B1 9 ⑲ B2 8 B3 5 B4 4 最も得点が高いものを取り出して、隣接する組合せを優先キューに加えていく 10 7 3 2 B1 9 ■ ⑯ B2 8 ⑱ B3 5 B4 4 10 7 3 2 B1 9 ■ ⑯ B2 8 ■ ⑮ B3 5 ⑮ B4 4 10 7 3 2 B1 9 ■ ■ ⑫ B2 8 ■ ⑮ B3 5 ⑮ B4 4 10 7 3 2 B1 9 ■ ■ ⑫ 等しい場合はどっちから取ってもいい B2 8 ■ ⑮ B3 5 ■ ⑫ B4 4 ⑭ 10 7 3 2 B1 9 ■ ■ ⑫ 5個取りだしたので終わり B2 8 ■ ■ ⑪ B3 5 ■ ⑫ B4 4 ⑭
必要最小限の組合せのみ、計算していることがわかる。
from heapq import heappop, heappush x, y, z, k = map(int, input().split()) aaa = sorted(map(int, input().split()), reverse=True) bbb = sorted(map(int, input().split()), reverse=True) ccc = sorted(map(int, input().split()), reverse=True) q = [(-(aaa[0] + bbb[0] + ccc[0]), 0, 0, 0)] ans = [] added = {(0, 0, 0)} for _ in [0] * k: s, ai, bi, ci = heappop(q) ans.append(-s) for da, db, dc in ((1, 0, 0), (0, 1, 0), (0, 0, 1)): aj, bj, cj = ai + da, bi + db, ci + dc if aj >= x or bj >= y or cj >= z: continue if (aj, bj, cj) in added: continue added.add((aj, bj, cj)) heappush(q, (-(aaa[aj] + bbb[bj] + ccc[cj]), aj, bj, cj)) print('\n'.join(map(str, ans)))
美味しい方から候補を広げていく
本番中に考えた方法。
$K$ 番以内に入る得点を作るのに用いられるケーキは、恐らくそこまで多くない。
この候補を適切に絞ることが出来れば、その中で全組み合わせを計算しても間に合う。
- $K$ 番以内に入る得点を作るのに用いられるケーキの候補を、それぞれのキャンドル毎に持つ
- 候補は、初期状態は各キャンドルの最大値のみであり、徐々に広げていく
- 候補中の全ての組み合わせの得点を保持し、候補を広げるたびに更新していく
- 新しい候補によって追加される得点の最大が、現時点の $K$ 番目以下になったらそこで打ち切る
採用すると得点の高くなるケーキとは、単純に得点の高いケーキでは無く、「同じキャンドルの中で最大の美味しさを持つケーキとの差が少ない順」である。
(サンプル2) X=3 Y=3 Z=3 K=5 A: 1 10 100 B: 2 20 200 C: 1 10 100
このような場合、まず各キャンドルの最大値(100,200,100)は確実に1番大きくなる。次に単独での美味しさが高いのは $B$ の20のケーキではあるが、それを選ぶということは、200を選ばないことを意味する。結果として合計は大きく減ってしまう。
以下のように最大値との差分を考えて、少ない方から順に採用するとよい。
A: -99 -90 0 B: -198 -180 0 C: -99 -90 0
まず、各キャンドルで最大の美味しさをそれぞれ $m_A,m_B,m_C$ とする。
各キャンドルで採用候補のケーキのリスト $D_A,D_B,D_C$ とし、それぞれ $m_A,m_B,m_C$ を追加しておく。
他のケーキは、同じキャンドルの中での最大値との差分をキーとして、全て一緒にしてソートする。どのキャンドル由来かは分かるようにしておく。このリストを $E$ とする。
現在の候補の全組み合わせから得られる得点のリストを $F$ とする。明らかな1番目として $m_A+m_B+m_C$ を追加しておく。
(1) $E$ から差分が少ないケーキを1つpopする。仮にキャンドル1のケーキだったとする(他のキャンドルでも考え方は同じ)。そのケーキの美味しさを $a$ とする。
このケーキを使って得られる最大得点は、$a+m_B+m_C$ である。 これは、現時点で候補として採用されていないケーキを1つ以上使って作れる最大得点である。
これが現時点の $F$ の $K$ 番目より大きければ、$a$ は最終的な $K$ 番以内の得点を作るのに必要である。 $a$ と $D_B,D_C$ 内の全ての組み合わせとの得点を $F$ に追加し、$K+1$ 番以下を切り捨てる。また、$D_A$ に $a$ を追加する。 (1) に戻って繰り返す。
そうでなく $a+m_B+m_C$ が $F$ の $K$ 番目以下なら、現候補内の組合せだけで $K$ 番目まで作れるので、そこで打ち切る。
毎回 $F$ をソートする必要があるが、繰り返し回数、$F$ のサイズ共にそこまで大きくならないので、十分高速に動作する。
x, y, z, k = map(int, input().split()) aaa = sorted(map(int, input().split()), reverse=True) bbb = sorted(map(int, input().split()), reverse=True) ccc = sorted(map(int, input().split()), reverse=True) ma, mb, mc = aaa[0], bbb[0], ccc[0] cakes = [(ma - a, 0) for a in aaa[1:]] cakes += [(mb - b, 1) for b in bbb[1:]] cakes += [(mc - c, 2) for c in ccc[1:]] cakes.sort() buf = [[ma], [mb], [mc]] ans = [ma + mb + mc] for d, i in cakes: if i == 0: a = aaa[len(buf[i])] if len(ans) == k and ans[-1] >= a + mb + mc: break buf[i].append(a) ans.extend(a + b + c for b in buf[1] for c in buf[2]) ans.sort(reverse=True) ans = ans[:k] elif i == 1: b = bbb[len(buf[i])] if len(ans) == k and ans[-1] >= ma + b + mc: break buf[i].append(b) ans.extend(a + b + c for a in buf[0] for c in buf[2]) ans.sort(reverse=True) ans = ans[:k] else: c = ccc[len(buf[i])] if len(ans) == k and ans[-1] >= ma + mb + c: break buf[i].append(c) ans.extend(a + b + c for a in buf[0] for b in buf[1]) ans.sort(reverse=True) ans = ans[:k] print('\n'.join(map(str, ans)))