めもめも

このブログに記載の内容は個人の見解であり、必ずしも所属組織の立場、戦略、意見を代表するものではありません。

E - Shiritori の解説

何の話かと言うと

atcoder.jp

この問題をネタに「トポロジカルオーダーで処理する」パターンの例を紹介します。

トポロジカルオーダーで処理する話については、まずは、下記の記事を参照してください。

enakai00.hatenablog.com

考え方

上記の記事では、「始状態」を順に削除していくという方法で処理を進めました。一方、今回の問題では、「終状態」を順に削除していく、という発想になります。(厳密にはちょっと違うのですが、詳しくはこの後の説明を読んでください。)

まず、単語 S1 の末尾3文字が単語 S2 の頭の3文字に一致する時に、S1 -> S2 の有効リンクを貼るという処理で単語の集合を有向グラフとみなします。

また、単語 n が自分の手番に回って来たとして、この時の状態価値関数を dp[n] とします。

・このままゲームを続けて確実に負ける場合 : dp[n] = -1
・確実に勝てる場合 : dp[n] = 1

たとえば、遷移先のない「終状態」の単語 n については、これが自分の手番に回ってくると負けが確定するので、dp[n] = -1 と確定します。

すると、単語 n に遷移できる単語 i については、dp[i] = 1 と確定します。(次に単語 n を選べば相手の負けが確定するので。)

この段階で、「終状態」、および、「終状態の1つ手前」の状態について、dp[n] が確定しました。

次に何を考えるかと言うと・・・

dp[i] = 1 である「終状態の1つ手前」の単語 i に関して、この単語 i は相手には渡せません。(これを渡すと負けが確定するから。)つまり、この単語 i は選択対象とはなりえず、はじめから存在しないのと同じことになります。そこで、単語 i に対する既存のリンクをすべて切ってしまいます。

この結果、遷移先が存在しない「終状態 n」が新たに発生して、これに対して dp[n] = -1 と確定します。すると、これの一歩手前の状態は dp[n] = 1 と確定して・・・・と、以下、同様の操作を続けていくことができます。


つまり、「終状態の一歩手前の状態へのリンクを切る」という操作を再帰的に繰り返すことで、勝敗が確定する状態をすべて辿ることができます。キュー(スタック)を使って再帰を実装するとこんな感じになります。

# link_to[n] : ノード n にリンクするノードのリスト
# link_from[n] : ノード n からリンクするノードのリスト
# link_from_count[n] : ノード n から伸びるリンクの数

  q = deque()
  for n in range(1, N+1):
    if link_from_count[n] == 0: # 遷移先の無い単語(終状態)
      q.append(n)

  while q:
    n = q.pop()
    dp[n] = -1 # 単語 n が回ってきたら負け
    for i in link_to[n]:
      if dp[i] != None: # 既に勝ちが決まっているものはスキップ
        continue
      dp[i] = 1 # 単語 n に回せる単語が回ってきたら勝ち
      for j in link_to[i]:
        link_from_count[j] -= 1 # i は選べないので i へのリンクを切る
        if link_from_count[j] == 0: # 新たな終状態が発生
          q.append(j)

先手が単語 n を選ぶとすれば、先手の勝敗は dp[n] の値で決まります。n は先手が後手に渡す単語になるので、dp[n] は後手の勝敗を示す点に注意してください。

  for n in range(1, N+1):
    if dp[n] == -1:
      print('Takahashi')
    elif dp[n] == 1:
      print('Aoki')
    else:
      print('Draw')

あとは、事前準備として、単語間の繋がりを見ながら、リンクのリストを構築する必要がありますが、愚直に二重ループすると TLE するので、ここはうまいことやってください。(最初、ここでTLEしているのに気づかずに、ちょっとはまっちゃいました。)

たとえば、頭、もしくは、末尾の3文字から対応する単語を検索できるディクショナリーを用意します。

  tail_dict = defaultdict(list)
  head_dict = defaultdict(list)

  for n in range(1, N+1):
    S[n] = f.readline().strip()
    head, tail = S[n][:3], S[n][-3:]
    tail_dict[tail].append(n)
    head_dict[head].append(n)
    for i in tail_dict[head]:
      link_from[i].append(n)
      link_to[n].append(i)
      link_from_count[i] += 1
    for i in head_dict[tail]:
      link_from[n].append(i)
      link_to[i].append(n)
      link_from_count[n] += 1

これでOKです。

atcoder.jp

D - Send More Money の解説

何の話かと言うと

atcoder.jp

この問題をネタに「しらみつぶし」系のコードの書き方のパターンを紹介します。

候補のリストをメンテナンスする

上記の問題、公式解説では、「0 〜 9 を重複なく割り当てられる文字は最大でも10種類なので、10! 通りを端から順に調べる」という方針が説明されていますが、単純なしらみつぶしではなく、N1 + N2 = N3 という条件を考慮しながら割り当て方法の候補を作ることで、さらに計算量を削減することができます。

具体的には、a 〜 z に対応する数字を一列ならべたリスト(対応する数字がない部分は None を入れておく)を「対応表」として、考え得る対応表を詰め込んだリスト candidate を用意します。はじめは、candidate は空です。

はじめに、S1 の 1 の位の文字を見て、その文字に 0 〜 9 を割り当てた 10 種類の対応表を candidate に詰め込みます。

次に、candidate 内のそれぞれの対応表に対して、S2 の 1 の位の文字を見て、その文字に 0 〜9 を割り当てて、対応表を更新します。重複する割り当てはスキップします。(この時点で 10*9 通りの対応表が candidate に詰まっています。)

続いて、candidate 内のそれぞれの対応表に対して、対応表の割り当てを用いて、N1+N2 の 1 の位を計算します。S3 の 1 の位の文字には、この値のみが割り当てられます。うまく割り当てができない場合は、その対応表は破棄します。

・・・という処理を 1 の位、10 の位、100 の位・・・と行っていきます。適切な割り当てができなくなった候補はどんどん破棄されるので、それぞれの位における候補の数は、200 〜 300 個程度に抑えられます。最終的に残った候補の 1 つを答えればOKです。

実際のコードはこちらになります。

atcoder.jp

S1, S2, S3 の桁数が異なる場合は、上の桁の足りない部分を仮想的な「0番目の文字」として処理しており、a 〜 z を 1 番目〜26番目の文字として、「0番目の文字」を含めた対応表を使っています。一番上の桁まで処理した時に、対応を埋めきれていない候補が残っている可能性もあるため、最後に残った候補については、1つずつ「検算」して、検算に通ったものを答えとして表示しています。

D - RGB Coloring 2 の解説

何の話かと言うと

atcoder.jp

上記の問題をネタに、(木構造とは限らない)単純無向グラフの基本的な取り扱いを説明します。

単純無向グラフは、多重リンクや自分から自分へのリンクを含まない無向グラフです。一般には、複数の連結成分に別れます。

データの持ち方

基本的には、各ノードに対して、リンクされたノードのリストを用意すればOKです。

  links = [[] for _ in range(N+1)] 
  for _ in range(M):
    a, b = list(map(int, f.readline().split()))
    links[a].append(b)
    links[b].append(a)

連結成分に対する処理

次のようにキューを用いると、あるノード n に対して、そこからすべての連結成分を辿ることができます。

    flag = [None] * (N+1) # 処理済みのノードを示す

    q = deque()
    flag[n] = True
    q.append(n)
    while q:
      i = q.pop()
      # ここでノード i についての処理をする
      for j in links[i]:
        if flag[j]:
          continue
        flag[j] = True
        q.append(j)

木構造の場合は、下記の記事で説明したように、親ノードの情報を与えることで「上から下にたどる」という順序を保証しました。

enakai00.hatenablog.com

今回は木構造とは限らないので、処理済みのノードを示すフラグを利用して、ループを避けています。木構造の場合も同じ手法を使うことができます。

連結部分ごとに個別に処理

連結とは限らないグラフに対して、連結部分ごとに処理をしたい場合は、次の様にすべてのノードについてループすればOKです。

  flag = [None] * (N+1) # 処理済みのノードを示す
  for n in range(1, N+1):
    if flag[n]:
      continue

    q = deque()
    flag[n] = True
    q.append(n)
    while q:
      i = q.pop()
      # ここでノード i についての処理をする
      for j in links[i]:
        if flag[j]:
          continue
        flag[j] = True
        q.append(j)

フラグのチェックによって、処理済みのノードはスキップされるので、同じ連結成分を2回処理することはありません。

問題の解法

で、最後に冒頭の問題の解法ですが、連結部分ごとに場合の数を計算して、それらの積で全体の答えになることから、上記のパターンが使えそうだと気がつきます。

それぞれの連結部分について、ノードを深さ優先探索でたどりながら、その時点で可能な配色を網羅していきます。

たとえば、

2 --- 1 --- 3
|-----------|

というループ状のグラフを 1, 2, 3 の順に処理する場合、

1. 可能な配色は、[R, None, None] (対称性から、G, B の場合はスキップして最後に3倍する。None は未定の意味)
2.可能な配色は、[R, G, None], [R, B, None]
3. 可能な配色は、[R, G, B], [R, B, G]

という様に順番に決まっていきます。あるノードを処理する際に、(深さ優先探索をしているので)そこからリンクされたノードの1つは確実に処理済みなので、追加する色の候補は、2種類以下になります。これにより、組み合わせの数を制限して、組合せ爆発を防いでいます。

具体的な実装としては、各ステップにおいて、その時点で可能性のある配色全体のリスト(candidate = [[R, G, None], [R, B, None]] みたいなやつ)をアップデートしていく感じになります。この配色リストをアップデートする処理を上記のパターンの

      # ここでノード i についての処理をする

という部分に当てはめればOKです。

ただし、この手のアップデートパターン(一種のDPですね)では、最初のノードに対する処理だけは別枠になるので、そのまま当てはめると場合分けの記述がごちゃごちゃします。そこで、まずは、深さ優先探索でグラフを辿る際のノード順のリストを作ります。

    tree_order = []
    q = deque()
    q.append(n)
    flag[n] = True
    while q:
      i = q.pop()
      tree_order.append(i)
      for j in links[i]:
        if flag[j]:
          continue
        flag[j] = True
        q.append(j)

この後、このリスト順にループを回せば、最初の要素 tree_order[0] に対する処理をきれいに分けて書くことができます。

完成版はこんな感じ。

import sys, copy
from collections import defaultdict, deque

def main(f):
  N, M = list(map(int, f.readline().split()))
  links = [[] for _ in range(N+1)] 
  for _ in range(M):
    a, b = list(map(int, f.readline().split()))
    links[a].append(b)
    links[b].append(a)

  flag = [None] * (N+1)
  result = 1
  for n in range(1, N+1):
    if flag[n]:
      continue

    tree_order = []
    q = deque()
    q.append(n)
    flag[n] = True
    while q:
      i = q.pop()
      tree_order.append(i)
      for j in links[i]:
        if flag[j]:
          continue
        flag[j] = True
        q.append(j)

    candidates = []
    i = tree_order[0]
    key = [None] * (N+1)
    key[i] = 0
    candidates.append(key)
    for i in tree_order[1:]:
      candidates_n = []
      for col_list in candidates:
        for c in range(3):
          ok = True
          for j in links[i]:
            if col_list[j] == c:
              ok = False
          if ok:
            col_list[i] = c
            candidates_n.append(copy.copy(col_list))
      candidates = candidates_n

    result *= len(candidates) * 3

  print(result)

main(sys.stdin)

atcoder.jp