情報オリンピック
JOI 2023/2024 2次予選 問題1 解答解説
A – カードゲーム 2 (Card Game 2)
この解説は本選(Aランク)を目指す生徒向けの内容です。一部省略している部分があります。
問題概要
与えられたリストの中から、特定の条件を満たす3つのカードの組み合わせが存在するかどうかを判定する問題です。
具体的には、ある値 a[i] に対して、a[i] + 3 と a[i] + 6 が両方ともリストに存在するかを調べます。
解法1: 愚直解(50点)
最も単純な方法は、3重ループで全ての組み合わせをチェックすることです。
n = int(input())
a = list(map(int, input().split()))
for i in range(n):
for j in range(n):
for k in range(n):
if a[i] + 3 == a[j] and a[i] + 6 == a[k]:
print("Yes")
exit()
print("No")
しかし、この解法は計算量が O(N³) となり、N=200,000 のような大きな入力ではタイムアウトしてしまいます。
よくある間違い
Pythonの in 演算子を使えば高速になると思いがちですが、リストに対する in は内部で線形探索を行っています。
# これは O(N²) のまま!
for v in a:
if v + 3 in a and v + 6 in a: # in はリストに対して O(N)
print("Yes")
exit()
リストに対する in 演算子は O(N) の計算量がかかるため、全体として O(N²) となり、やはり大きな入力では間に合いません。
解法2: 二分探索(100点)
データをソートして二分探索を使うことで、探索を O(log N) に改善できます。
import bisect
n = int(input())
a = list(map(int, input().split()))
a.sort()
for v in a:
# 二分探索で v+3 と v+6 を探す
idx3 = bisect.bisect_left(a, v + 3)
idx6 = bisect.bisect_left(a, v + 6)
if idx3 < n and a[idx3] == v + 3 and idx6 < n and a[idx6] == v + 6:
print("Yes")
exit()
print("No")
この解法の計算量は O(N log N) です。
解法3: set型を使う(100点・推奨)
最も簡潔で効率的な解法は、リストを set(集合)型 に変換することです。
n = int(input())
a = set(map(int, input().split()))
for v in a:
if v + 3 in a and v + 6 in a:
print("Yes")
exit()
print("No")
set型に対する in 演算子は O(1)(平均)で動作するため、全体として O(N) で解くことができます。
ポイント: 「要素の順序や重複が重要でない」かつ「要素の存在確認を頻繁に行う」場合は、set型を使うことで大幅に高速化できます。
計算量の比較
| 解法 | 計算量 | 得点 |
|---|---|---|
| 3重ループ | O(N³) | 50点 |
| リスト + in | O(N²) | 50点 |
| 二分探索 | O(N log N) | 100点 |
| set型 | O(N) | 100点 |