【AtCoder】PythonでABC236 D問題(Dance)を解く

AtCoder青色のkym3535と申します。昨日のABC236ではD問題が難しく、C問題の正解者数6900人に対してD問題1800人と、かなりの崖が出来ていました。私自身も結構時間をとられてしまいましたが、無事正解できましたので、解説記事を書きたいと思います。

問題

(AtCoder Beginner Contest 236) D - Dance

$1, 2, \cdots, 2N$と番号づけられた$2N$人の人が舞踏会に参加します。彼らは$N$個の$2$人組に分かれてダンスを踊ります。

$2$人組を構成する人のうち、番号の小さい方の人が人$i$、番号の大きい方の人が人$j$のとき、その2人組の「相性」は$A_{ij}$です。
$N$個の$2$人組の相性がそれぞれ$B_1$,$B_2$,...,$B_N$であるとき、「舞踏会全体の楽しさ」はそれらのビットごとの排他的論理和である$B_1 \oplus B_2 \oplus \cdots \oplus B_N$ です。

「2N人の参加者がN個の2人組に分かれる方法」を自由に選べるとき、「舞踏会全体の楽しさ」としてあり得る最大値を出力してください。

考察

ぱっと見て、単純に和をとるのでは無くXORをとるため、計算をうまく工夫してオーダーを減らすのは難しそうと感じた。
全探索するのかと思って制約を確認すると、$N \leq 8$が目に入ったので、$8!$が(競プロ的には)大したことないことと照らし合わせて、詳しく考えてみることにした。

たとえば6人の場合(N=3)、人1とペアになる人は2,3,4,5,6の5通り。
・仮に(1,2)をペアにした場合、人3とペアになる人は4,5,6の3通り。
・仮に(1,3)をペアにした場合、人2とペアになる人は4,5,6の3通り。
・仮に(1,4)をペアにした場合、人2とペアになる人は3,5,6の3通り。
・(以下略)
のように、「残っている人の中で最も番号が若い人と、他の誰かをペアにする」を繰り返していくと重複無く数え上げられる。
6人の場合の組み合わせは、
・最初の選択肢が5通り
　・次の選択肢が3通り
　　・最後の選択肢が1通り(ちょうど2人が残っている)
なので、$5\times 3 \times 1 = 15$ 通りと計算できる。

同様に、最大16人の場合を考えると、すべての組み合わせは $15\times 13\times 11\times \cdots \times 1 $通り。サクっとPythonのコンソールで計算してみる。

約200万通り。各パターンについて、8回XORをとった後、最大値の更新処理をするだけなので、全探索で十分間に合いそうだ。

実装

実装方針の決定

考察まではスムーズに進んだが、実装が難しかった。

単純にforループで回すのは難しく、再帰を使う必要がありそうなのは薄々感じ取れた。
再帰を使う際には、再帰したい関数の仕様をきちっと決める必要がある。入力として何を受け取り、何を返すかだ。

しばらく考えたが、次のような方針で書くことにした。
関数は、残っている人の番号のリストを入力として受け取る。そのリストの人たちを2人ずつペアにしたすべての場合について、$A_{ij}$のXORをとった値を計算し、それらのリストを出力として返す。そうすると、($A_{ij}$の値を適当に設定しますが、)

1. $[0,1,2,3,4,5]$について解こう。
2. 人0と1をペアにした場合、残りの$[2,3,4,5]$を2人ずつに分けた3通りで、それぞれ$A_{ij}$のXORはどうなるだろうか?再帰関数に任せよう!
3. 再帰した結果、$A_{2,3} \oplus A_{4,5}=10$ ([(2,3),(4,5)]のペアにした場合)、$A_{2,4} \oplus A_{3,5}=15$ ([(2,4),(3,5)]のペアにした場合)、$A_{2,5} \oplus A_{3,4}=8$ ([(2,5),(3,4)]のペアにした場合)、だったので、$[10, 15, 8]$ が返ってきた。
4. $A_{0,1}=1$ なので、それぞれXORをとると　$10\oplus1=11$、　$15\oplus1=14$、　$8\oplus1=9$　なので、$[11, 14, 9]$ だな。
5. 次に、人0と2をペアにした場合、残りの$[1,3,4,5]$について、同様に再帰関数に任せよう!
6. 再帰した結果、$[10, 14, 9]$ が返ってきた。
7. $A_{0,2}=16$なので、それぞれXORをとると$[26, 30, 25]$ だな。これも答えとしてあり得るな。
8. 0と3、0と4、0と5をペアにする場合も、同様に計算しよう。
9. 最終的に、答えとしてあり得る値として、15通りの数値が得られる。これらの数値のリストを、出力として返そう。

のように、再帰の前後でうまくつながってくれそうだ。
入力のリストの長さが2の場合は、再帰を止める。このときの返り値は、入力された2人の相性の数値が含まれたリストだ。

リストを再帰するごとに返すのは明らかに効率が悪いが、これ以上いいやり方も思い浮かばないのでとりあえず書いてみることにした。

最大ケースの確認

ひととおり書き上げて、サンプルケースはすべて通ることを確認した。しかし、実装効率が悪いため最大ケースで通るか不安だ。
幸い、入力の規模は小さくて済むので、急いで適当な最大ケースを作り、コードテストを回してみる。

maxcase.txt

8
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
1 2 3 4 5 6 7 8 9 10 11 12 13 14
1 2 3 4 5 6 7 8 9 10 11 12 13
1 2 3 4 5 6 7 8 9 10 11 12
1 2 3 4 5 6 7 8 9 10 11
1 2 3 4 5 6 7 8 9 10
1 2 3 4 5 6 7 8 9
1 2 3 4 5 6 7 8
1 2 3 4 5 6 7
1 2 3 4 5 6
1 2 3 4 5
1 2 3 4
1 2 3
1 2
1

PypyとPythonの選択だが、Pypyは基本的にPythonより3倍程度速いが、再帰はかなり遅い特徴がある。
今回の問題は再帰を使っているので、Pypyの苦手分野だが、再帰の段数はたかだか8回なので、Pypyでもそこまで遅くはならないだろうと考え、まずはPypyに投げてみる。

約800(msec)。実際はXORにもっと複雑な数値が入るが、実行時間の大半はループ・再帰処理だろうと考えられるので、余裕で間に合うだろうと判断し、そのままPypyに投げることにした。

提出

清書した解答【提出 #28780096】
本番での解答【提出 #28743245】

本番での解答は、かなり汚いところもありますがご容赦ください…。

ABC236D_seisho.py

import sys
def input():
    return sys.stdin.readline()[:-1]

a = []

def f(lis):
    """
    lis の中から2個とって、残りの要素について再帰する。
    lis の中味を2個ずつ分けたすべての組み合わせについて、
    a[i][j]のxorをとった値のリストを返す。
    
    ◆入力
    - lis : リスト。長さは偶数にすること。  例 : [0,1,4,5]
    ◆出力
    - sub_ans_list  : lisを2つずつ組にしたすべての場合について、A_ijのxorをとった値のリスト。
                        例 : [ a[0][1]^a[4][5], a[0][4]^a[1][5], a[0][5]^a[1][4] ]
    """
    ans_list = []
    
    # リスト先頭の人の番号
    pi = lis[0]
    for j in range(1,len(lis)):
        # リストj番目の人の番号
        pj = lis[j]
        # 先頭の人とj番目の人の相性
        bij = a[pi][pj]
        
        if len(lis)>2:
            # lisの中から、先頭とj番目を取り除いた小リストを作る。
            sub_list = lis[1:j]+lis[j+1:]
            
            # 小リストについて問題を解き、あり得るすべてのxorの値のリストを得る。
            sub_ans_list = f(sub_list)
            
            for sa in sub_ans_list:
                ans_list.append(sa^bij)
        else:
            ans_list.append(bij)
    
    return ans_list

def main():
    n = int(input())

    # aを行列の形に整形する。
    # 人xとyの相性は、(xとyを0-based indexで表記すると)
    # a[x][y]と表される。
    for i in range(2*n-1):
        a_row = list(map(int, input().split()))
        a_row = [0]*(i+1) + a_row
        a.append(a_row)

    # lis = [0,1,2, ... , 15]   # n=8の場合
    lis = list(range(2*n))

    # 問題を解く
    ans_list2 = f(lis)
    print(max(ans_list2))
    
if __name__ == '__main__':
    main()

後日談：ジェネレータを用いた実装

翌日の仕事の休憩時間、ふと、ジェネレータを使うとすべてのペアの組み合わせを綺麗に探索できることに気づいた。
再帰関数でyield文を使ったコードなんて書いたことが無かったが、試してみると1発で通ってしまった。
実行時間は約1200(msec)で、先のコードよりも遅かったが、参考のため、こちらも載せておきます。

別解(ジェネレータ)【提出 #28777915】

ABC236D_seisho.py

import sys
def input():
    return sys.stdin.readline()[:-1]

def pairs(lis):
    """
    lis内の要素を2個ずつのタプルにまとめたリストを、
    すべての分け方について順番に返す。
    
    例 : for pattern in pairs([10,20,30,40]):
        (略)
    ===> pattern には、順に
            [(10,20),(30,40)]
            [(10,30),(20,40)]
            [(10,40),(20,30)]
        が代入され、for文が繰り返される
    
    """
    if len(lis) == 2:
        yield [ tuple(lis) ]
    else:
        for i in range(1,len(lis)):
            rest = lis[1:i] + lis[i+1:]
            new_tup = (lis[0], lis[i])
            for sub_pair in pairs(rest):
                # 子のジェネレータから1つずつ返ってきているペアのリストは、
                # 新しい要素を追加して yield さえしてしまえば、自身は2度と
                # 触る必要が無いので、append して書き換えても問題ない。
                sub_pair.append(new_tup)
                yield sub_pair
def main():
    n = int(input())

    # aを行列の形に整形する。
    # 人xとyの相性は、(xとyを0-based indexで表記すると)
    # a[x][y]と表される。
    a = []
    for i in range(2*n-1):
        a_row = list(map(int, input().split()))
        a_row = [0]*(i+1) + a_row
        a.append(a_row.copy())

    # lis = [0,1,2, ... , 15]   # n=8の場合
    lis = list(range(2*n))

    ans = 0
    for pair_list in pairs(lis):
        c = 0
        for x,y in pair_list:
            c ^= a[x][y]
        if c>ans:
            ans = c
    print(ans)
    
if __name__ == '__main__':
    main()