情報オリンピック

JOI 2023/2024 2次予選 問題1 解答解説

A – カードゲーム 2 (Card Game 2)

この解説は本選(Aランク)を目指す生徒向けの内容です。一部省略している部分があります。

問題概要

与えられたリストの中から、特定の条件を満たす3つのカードの組み合わせが存在するかどうかを判定する問題です。 具体的には、ある値 a[i] に対して、a[i] + 3a[i] + 6 が両方ともリストに存在するかを調べます。

AtCoderで問題を見る →

解法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点
リスト + inO(N²)50点
二分探索O(N log N)100点
set型O(N)100点