ABC200のD問題 Happy Birthday! 2 を工夫して解く - 競技プログラミング悪戦苦闘日記

残念ながらABC200のD問題を解けなかったので、解説しつつふりかえります。

問題の概要

N 個の正整数からなる数列 $A = (A_1, A_2, \dots, A_N)$ が与えられます。以下の条件を全て満たす 2 つの数列
$B = (B_1, B_2, \dots, B_x), C=(C_1, C_2, \dots, C_y)$ が存在するか判定し、存在する場合はひとつ出力してください。
（省略）
$(A_{B_1} + A_{B_2} + \dots + A_{B_x})$ を200で割ったあまりと、 $(A_{C_1} + A_{C_2} + \dots + A_{C_y})$ を200で割ったあまりが等しい

$A = (A_1, A_2, \dots, A_N)$ の中から200で割った余りが等しくなる足し算の組み合わせが2組以上あるかどうか判定して、あればそれを出力、なければ'No'と出力する問題です。

難しいポイント

Nが最大で200なのですが、足し算の組み合わせを単純に考えると $2^{200}$ となり、とてつもない数になります。
何かに気づいて間引かないといけないです。

解き方の方針

公式解説の「鳩の巣原理」は気づかなかったです。。。
余りのパターンとしては200で割ったあまりなので、0から199の200通りしかありません。
なので少なくとも201個合計のパターンあれば、2組のパターンが存在してそれが答えになります。
それを考えると、ビット全探索で $2^{200}$ までやる必要はなく、数列の中に同じ数字はないので $2^{8} = 256$ まで全探索すれば必ず答えはあるはずです。
なければ'No'と出力します。

実装

2次元配列で、同じ余りになる組み合わせを記録しました。
例えば、ans_list[5] = [1, 2, 5] の時は、a_arr[0]+a_arr[1]+a_arr[4]の余りが5になるように記録します。

N = int(input())
a_arr = list(map(int, input().split()))
# 2**8の中で答えの組み合わせがある。Nが8より小さければ全通り探索する
n = min(N, 8)
# 同じ余りとなる組み合わせのリストを保存するリスト
# ans_list[5] = [1, 2, 5] はa_arr[0]+a_arr[1]+a_arr[4]の余りが5になる
ans_list = [[] for _ in range(200)]
for bit in range(2**n):
    tmp_sum = 0
    tmp_list = []
    for i in range(n):
        if bit & (1 << i):
            tmp_sum += a_arr[i]
            tmp_list.append(i + 1)
    amari = tmp_sum % 200
    # 同じ余りの組み合わせがあれば答えを出力して終了する
    if len(ans_list[amari]):
        print('Yes')
        print(len(ans_list[amari]), *ans_list[amari])
        print(len(tmp_list), *tmp_list)
        exit()
    ans_list[amari] = tmp_list

print('No')