JOI 2023/2024 2次予選問題1 解答解説 | 秀英塾コラム

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

対象外：1次予選で4問完答できない人　または　2次予選突破する実力がすでにある人

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

与えられたリストの中から、特定の条件を満たす3つのカードの組み合わせが存在するかどうかを判定する問題です。

愚直解（50点）

この手の問題は、まず全探索することを検討してください。今回は３枚のカードの組み合わせが条件を満たすかを調べたいので、３重ループを行います。

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")

このプログラムは、小課題１と小課題３で得点でき、50点が得られます。 3重ループだとO(n^3)なので、N=200,000のケースで8×10^15となりTLEになります。

よくある間違い

forループの入れ子（ネスト）を取り除くため、以下のようなプログラムを書いている人が見受けられましたが、 listに対してinを使って検索するのは、内部的にはforループを使って前から順番に探す処理（線形探索）が行われるので、ループ１つあたりの計算量がO(n)となり全体としてO(n^2)で高速化できません。

n = int(input())
a = list(map(int, input().split()))

# これは O(N²) のまま！
for v in a:
    if v + 3 in a and v + 6 in a:  # in はリストに対して O(N)
        print("Yes")
        exit()

満点解法１: 二分探索（100点）

N=200,000ということは、逆に考えるとO(n log n)であるソートは間に合います。ソートができれば2分探索ができるので、以下のプログラムで満点回答が狙えます。

n = int(input())
a = list(map(int, input().split()))
a.sort()

for i in range(n):
    # x + 3を探す
    ok = -1
    ng = n

    while ng - ok > 1:
        mid = (ok + ng) // 2
        if a[mid] <= a[i] + 3:
            ok = mid
        else:
            ng = mid
    if a[ok] != a[i] + 3:
        continue

    # x + 6を探す
    ok = -1
    ng = n

    while ng - ok > 1:
        mid = (ok + ng) // 2
        if a[mid] <= a[i] + 6:
            ok = mid
        else:
            ng = mid
    if a[ok] != a[i] + 6:
        continue

    print("Yes")
    exit()

print("No")

この解法の計算量は O(N log N) です。 2分探索は慣れるまでが大変ですが、コツをつかむと色々な問題で使えますのでぜひ習得してください。

満点解法２: 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")

2つ目のプログラムのlistにしていたところをsetに変えるだけで満点回答になります。

setでは、要素の順番や重複した要素を管理できない代わりに、「指定した要素があるかないか」をO(1)で調べることができます。

n個の要素それぞれにO(1)で調べれば良いので、プログラム全体の計算量はO(n)となります。

今回の問題のように、カードに書かれた数字の順番の情報や同じ番号のカードの情報が必要ないときは積極的にsetを使ってください。

ポイント: 「要素の順序や重複が重要でない」かつ「要素の存在確認を頻繁に行う」場合は、set型を使うことで大幅に高速化できます。

計算量の比較

解法	計算量	得点
3重ループ	O(N³)	50点
リスト + in	O(N²)	50点
二分探索	O(N log N)	100点
set型	O(N)	100点