043 - Maze Challenge with Lack of Sleep（★4）の解説

何の話かと言うと

atcoder.jp

この問題をネタに「重み0のリンクで結合されたノードを同一視する」という話をします。

一般的な解法

まず、この問題の一般的な解法は、01-BFS になります。

方向転換を考慮する必要があるので、座標 (x, y) + 方向（上下方向 or 左右方向）をノードとします。たとえば、

・(1, 1, '|') : 座標 (1, 1) を上下方向に移動中
・(1, 1, '-') : 座標 (1, 1) を左右方向に移動中

という感じです。そして、これらのノードをリンクで結合します。

座標を変えずに方向だけ変えるリンクを重み 1 とします。

・(1, 1, '|') <----> (1, 1, '-') : 重み 1

ただし、スタートとゴール地点での方向転換は重み 0 にします。（スタート地点での方向転換は自由なので。）

今向いている方向への移動を表すリンクを重み 0 とします。

・(1, 1, '|') <----> (1, 2, '|') : 重み 0
・(1, 1, '-') <----> (2, 1, '-') : 重み 0

あとは、これらで構成される無向グラフについて、スタートからゴールの最小距離を検索します。（スタートとゴールでの方向は任意の方向を1つ決めればOK。スタートとゴール地点の方向転換は重み 0 に注意してください。）

一般には、ダイクストラを使えばOKですが、Python だとギリギリ TLE になります。

そこで、ダイクストラより効率的な 01-BFS を利用します。（これは、リンクの重みが 0 / 1 だけの場合に使えます。）

※ 01-BFS の解説はこちらがお勧め。
betrue12.hateblo.jp

で、実際の結果はこうなります。

atcoder.jp

いちおうパスしていますが、1838 ms で、けっこうギリギリ感はあります。

重み 0 でつながるノードを同一視

上記のコードでは、conv() 関数で、「座標 + 方向（0 / 1 で表現）」で表した状態を単一の数値のノード番号に変換しています。

仮に、(1, 1, 0) と (2, 1, 0) が重み 0 のリンクで繋がっていた場合、これらを区別せずに同一のノード番号にマッピングしてしまっても構いません。

ただし・・・・

それぞれのノードがリンクを持っている場合、それぞれのリンクを融合するという、ちょっと面倒な処理が必要になります。

なのですが・・・

上記のコードでは、リンク情報を構築する際に、

  for y in range(1, H+1):
    for x in range(1, W+1):

という方向にマス目をスキャンしながら、

・(x, y) から見て、右と下にあるマス目 (x+1, y)、および、(x, y+1) とのリンクを追加する

という処理をしています。（この順番であれば、左と上のマス目は、ダブルカウントになるので、考えなくてよいですよね・・）

そうすると・・・

・(x, y) から見て、右と下にあるマス目 (x+1, y)、および、(x, y+1) は、スキャン前なので、まだリンク情報を持っていません。

したがって、「(x+1, y)、もしくは、(x, y+1) を (x, y) と同一視する」という処理、つまり、「(x+1, y)、もしくは、(x, y+1) を (x, y) と同一のノード番号にマッピングする」という処理は、リンクの引き継ぎを気にせずに行うことができます。

ノードの同一視については、（Union-Find の発想を借りて）次のように実装します。

まず、conv() では、ノード番号を取得した後、さらにその親のノード番号を返すようにします。

def conv(x, y, d):
  global H, W, parent
  i = x + (y-1) * W + d * (H*W)
  if parent[i] == None:
    parent[i] = i
  return parent[i]

で、ノードを同一視させたくなったところで、親の情報を書き換えます。

      if exists(x+1, y):
        p0 = conv(x, y, 0)
        p1 = conv(x+1, y, 0)
        parent[p1] = p0

一般には、親の親を再帰的に検索する必要がありますが、今回の場合は、スキャンの順序を（落ち着いて）考えると、再帰的な検索は不要とわかります。

実際にこれを実装したのがこちらになります。

atcoder.jp

801 ms で、確かに高速化されました。（ちなみに、この処理をしておけば、ダイクストラでもギリギリ間に合いました。）

とことんやってみる

上記のコードでは、隣り合うマス目を（連鎖的に）同一視していますが、よく考えると、スタート地点とゴール地点の方向転換についても、同一視ができます。（スタートとゴールでは方向転換のリンクは重み 0 なので。）

ただし、この部分では、先に触れたノードの引き継ぎが必要になります。さらに、同一視の連鎖が一方向にならない可能性があるので、（本物の Union-Find と同様に）親の検索を再帰的に実装する必要があります。

def conv(x, y, d):
  global H, W, parent
  i = x + (y-1) * W + d * (H*W)
  if parent[i] == None:
    parent[i] = i
  while i != parent[i]:
    i = parent[i]
  return i

ここまでやる必要があるかは疑問ですが、仮に、ここまでやると、重み 0 のリンクが完全になくなり、重み 1 のリンクのみになります。なので、01-BFS ではなく、普通の BFS での探索が可能になります。

というわけで、ものは試しでやってみました。

atcoder.jp

965 ms とちょっと遅くなりました。親の再帰的な検索の時間がオーバーヘッドになっているような気がします。（たぶん。）

040 - Get More Money（★7）の解説

何の話かと言うと

atcoder.jp

この問題、公式解説では、最小カット問題（＝最大フロー問題）に帰着して、Dinic 法で解くという解法が紹介されていますが、もうちょっと素朴な「ナップサック問題」としても解ける（解けた）ので紹介します。

ナップサック問題との類似性

まず、n 番目の家には、n+1 番目以降の家の鍵が置いてあることから、家を訪れる順番は、基本的に番号順になります。つまり、n=1,2,...,N について、それぞれの家に入る／入らないという判断をしていきます。

また、n 番目の家を入る／入らないという判断をするごとに、この後に入ることのできる家に対する制約条件が変化していきます。

これは、ナップサック問題に類似した状況です。

ナップサック問題では、n=1,2,...,N 個の荷物のそれぞれについて、入れる／入れないという判断をしていきますが、判断をするごとにナップサックの残り容量が減って、追加で入れられる荷物に対する制約が変化していきます。

このようなナップサック問題では、残り容量をキー、その時点での価値をバリューとするディクショナリーに情報を記録していくのが効率的です。

今回の場合は、未収集の家の鍵が徐々に減っていくので、たとえば、「現在残っている鍵の状態」をキーにすることが考えられます。最大 N 種類の鍵があるので、それぞれについての残り数を並べたタプルをキーにしてみましょう。初期状態は、次のように構成します。

import sys, copy

def main(f):
  global N, W, A, C
  N, W = map(int, f.readline().split())
  A = [None] + list(map(int, f.readline().split()))
  C = [[] for _ in range(N+1)]
  num_keys = [0] * N
  for n in range(1, N+1):
    C[n] = list(map(int, f.readline().split()))
    C[n] = C[n][1:]
    for c in C[n]:
      num_keys[c-1] += 1
  num_keys = tuple(num_keys)

  dp = defaultdict(lambda: -10**30)
  dp[num_keys] = 0

  print(dp)

#main(sys.stdin)

# input
5 5
5 2 10 3 6
1 3
1 3
0
1 5
0
# output
{(0, 0, 2, 0, 1): 0}

この例では、3番目の家の鍵が2本、5番目の家の鍵が1本残っており、現在の価値（）は 0 です。

この後は、n=1,2,...,N のループを回しながら、ディクショナリー dp のキーを取り出して、その状態から n 番目の家に入る場合、入らない場合の状態を追加していきます。家に入った場合は、その家にある鍵を減らしていきます。もちろん、残りの鍵が0でない家には入れません。

で、これを（効率はあまり考えずに）素朴に実装して提出するとこうなります。

atcoder.jp

AC x 10, TLE x 2 で、予想通り（？）TLE していますが、max ケースでも AC しているものもあり、方向性としては悪くなさそうです。

状態数の削減

この手のナップサック問題の場合、n が増えるごとに状態が倍になるので、状態数が爆発的に増えることにより、状態（ディクショナリーのキー）に関するループが長くなります。なんらかの方法で、状態数を減らす必要があります。

例えば、普通のナップサック問題であれば、ナップサックの総容量 W とすると、状態は 0 〜 W の整数に限定されるので、W が小さければ、自然と状態数も少なくなります。（原理的には、状態数は倍々で増えますが、結果的に残り容量の値が被る場合が多発します。）

では、今回の場合、何らかの方法で「状態」の範囲を限定して、状態数を減らすことができないでしょうか？

まず、「鍵の残り数」で考えることをやめます。重要なのは、残り数が0か0でないか（ある家に入れるか入れないか）、という事で、0 でない場合の具体的な値は本質ではありません。

・・・でも、残り数をカウントしないと、途中経過がトラッキングできないのでは・・・・？

そうでもありません。たとえば、家1に家3の鍵がある場合、家1に入らなければ、その時点で、家3に入れないことは確定します。（家3の鍵の残り数がいくらかは考える必要がありません。）つまり、それぞれの家に「入れる／入れない」という0/1のフラグをN個持っておき、ある家に入らないという選択をするごとに、入らなかった家にある鍵の家のフラグを1にすればよいのです。こうすれば、 の整数値をディクショナリーのキーにすることができて、状態数を減らすと共に、（重い）タプルの操作を無くすことができます。

・・・いやそれでも、 ってめちゃめちゃ大きくないです？

確かに・・・・。

いや、まだ諦める必要はありません。ここがポイントなのですが、今回の場合、それぞれの家にどの鍵があるかは（実際に入って確かめなくても）事前にすべてわかっています。「家1に家3の鍵がある場合、家1に入らなければ、その時点で、家3に入れないことは確定します」と言いましたが、家3に入れないということは、家3に鍵がある家にも入れない、という事実が確定します。つまり、1つの家に入らないと、その先で入れなくなる家が連鎖的に増えていきます。こう考えると、状態フラグは必然的に1が多くなり、実際に発生する状態の数は、 よりもずっと小さくなる可能性があります。

さらに、n 番目の家の処理がおわれば、この後は、n+1 〜 N 番目のフラグのみが必要で、それ以前のフラグの値は参照する必要がありません。そこで、n 番目を家の処理を終えたら、n 番目のフラグは強制的に 0 にします。こうすれば、処理がすすむごとに状態フラグの（実質的な）桁数が減っていき、状態数もどんどん減っていきます。

このあたりの処理をまとめて実装するとこんな感じです。連鎖的に入れなくなる家は、事前に計算して、bitmask として用意しておきます。

  mask = [0] * (N+1) # mask[n] : 家 n に入らなかった時に（連鎖的に）入れなくなる家を示す bitmask
  for c in C[N-1]:
    mask[N-1] |= 1<<(c-1)
  for n in range(N-2, 0, -1):
    for c in C[n]:
      mask[n] |= 1<<(c-1)
      mask[n] |= mask[c]

  st = 0
  dp = defaultdict(lambda: -10**30)
  dp[st] = 0
  for n in range(1, N+1):
    dp_n = defaultdict(lambda: -10**30)
    for st in dp.keys():
      # 家 n に入らない場合
      st_n = st
      st_n |= mask[n] # 連鎖的に入れなくなる家のフラグを立てる
      st_n ^ 1<<(n-1) # n 番目のフラグを強制的に 0 にする
      dp_n[st_n] = max(dp_n[st_n], dp[st])

      if st & 1<<(n-1) == 0: # 家 n に入る場合
        r = A[n] - W
        dp_n[st] = max(dp_n[st], dp[st] + r)
    dp = dp_n

で、これを提出すると・・・

atcoder.jp

惜しいです。TLE が 1 つだけ残ります。

状態の刈り込み

最後の手段で、状態の刈り込みを行いましょう。

ある2つの状態 A, B を比べて、

・A の方が B よりも状況的に不利で、かつ、その時点の価値も B より低い

という状況があった場合、状態 A をそれ以上トラッキングする必要はありません。その状態はディクショナリーから削除できます。

今の場合、

・B で入れない家は、すべて、A でも入れない

という場合、A は B よりも状況的に不利です。ビット表現であれば、flag_A | flag_B == flag_A と表せますね。このような組 A, B について、さらに dp[A] <= dp[B] であれば、A を捨てることができます。

ただ、今回の場合、このチェックを漏れなく行うには、状態数を として、 のループが必要になります。 が大きいとこの処理自体が遅くなる可能性もあります。

で、ここはぶっちゃけ状況しだいなのですが、仮に、この刈り込みが初期段階から強く効いたとすると、刈り込みによって状態数 が抑えられて、刈り込みの処理自体も遅くならずに済む可能性もあります。とにかくやってみましょう。

atcoder.jp

無事にすべて通りました。

039 - Tree Distance（★5）の解説

何の話かと言うと

atcoder.jp

この問題をネタに、「頂点についてのループと辺についてのループ」の話をします。・・・・というのは、ちょっと無理矢理感があって、ぶっちゃけは、別解を紹介したいだけです。

辺についてのループ

一般に、グラフの問題は、「N 個の頂点について順番に処理する」という発想で解きますが、この問題に関して言うと、（ちょっとした発想の転換で）「N-1 個の辺について順番に処理する」という方法で、驚くほど簡単に解くことができます。詳しくは、公式解説を参照ください。

頂点についてのループでは解けないの？

なのですが、「頂点ごとに考える」という普通の発想でも解けなくはないので、別解として、こちらの解法を紹介します。

この問題では、任意の2頂点の全ての組みについて、それらの最短経路を考える必要がありますが、頂点の数（）を考えると、各頂点を1回ずつしか処理できそうにありません。それでは、たとえば、深さ優先探索で、木を下から順番にたどった時に、ある特定の頂点について、どのような経路をカバーすることができるでしょうか？

一般に、木構造グラフでは、任意の2点を結ぶ最短経路は、「（共通の親まで）登って、そこからまた下る」という形になりますが、少なくとも「親まで登る経路の距離」は、深さ優先探索のDPでまとめあげることができます。つまり、

・dp[i] = "頂点 i のすべての子孫 j についての「j から i に至る距離」の合計"

が計算できます。もうちょっと厳密に言うと、「i を頂点とするサブツリーのノード数」の情報が補助的に必要になるので、

・dp[i] = (c, d) : i を頂点とするサブツリーについての (ノード数, 「子孫 j から親 i の距離」の合計)

という形になります。

import sys
from collections import defaultdict, deque

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

# dp[i] = (c, d) : i を頂点とするサブツリーについての (ノード数, 「子孫 j から親 i の距離」の合計)
  dp = [(0, 0)] * (N+1)

  q = deque()
  q.append((1, 0)) # (node, parent)
  while q:
    i, p = q.pop()
    if i < 0:
      c, d = 0, 0
      for j in link[-i]:
        if j == p:
          continue
        c_j, d_j = dp[j]
        c += c_j
        d += d_j + 1*(c_j) # j->i の距離 1 を各子孫について加える
      dp[-i] = (c+1, d)
      continue

    q.append((-i, p)) # 帰りがけ順で処理するための番兵

    for j in link[i]:
      if j == p:
        continue
      q.append((j, i))

  print(dp[1:])
main(sys.stdin)

# input
4
1 2
2 3
1 4

# output
[(4, 4), (2, 1), (1, 0), (1, 0)]

得られた距離をすべて足しあげれば、「登って終わり」というパターンの経路についての合計が得られます。

「登ってから下る」パターンはどうするの？

もちろん、このままでは、まだ、「すべての頂点の組み合わせ」がカバーできていません。頂点 i に注目した際に、

・頂点 i まで登って終わり

という経路が網羅されている状況ですので、これに加えて、

・頂点 i の直下の子ノード j < k について、「サブツリー j 内のノードから i まで登って、サブツリー k 内のノードまで下る」

という経路を加える必要があります。これをすべての頂点 i について網羅すれば、あらゆる経路がカバーされますよね・・・？

　

（i を折り返し点とするものは、左右の対称性があるので、一方のみを採用するとしてください。）

で、ここが頭の使い所なのですが・・・

実は、すでに計算ずみの dp[i] を利用すると、「i を折り返し点とする全経路の距離の合計」が簡単に計算できます。

今、頂点 i の下に直接ぶら下がる 2 つの頂点 j, k について、

・j の子孫 → i → k の子孫

という経路について、前半と後半を分けて考えます。

・前半：j の子孫 → i
・後半：i → k の子孫

ここで、

・(c_j, d_j) = dp[j]

と置くと、k の子孫を 1 つ固定した場合、これに至る経路は c_j（j 以下のサブツリーのノード数）本ありますが、これらの前半の経路の総距離は、d_j + 1 * c_j になります。（j の子孫 -> j の各経路について、j->i の距離 1 を加えている点に注意してください。）

したがって、c_k 個ある（k 自身を含めた）k のすべての子孫を考慮すると、前半の経路の総計は、

・(d_j + 1 * c_j) * c_k

で計算されます。

同様に、

・(c_k, d_k) = dp[k]

と置いて、j の子孫を 1 つ固定した場合、ここからスタートする経路は c_k（k 以下のサブツリーのノード数）本あり、これらの後半の経路の総距離は、d_k + 1 * c_k になります。したがって、c_j 個ある（j 自身を含めた）j のすべての子孫を考慮すると、後半の経路の総計は、

・c_j * (d_k + 1 * c_k)

で計算されます。

前半と後半を合わせると、i で折り返す経路の総距離は、

・(d_j + c_j) * c_k + c_j * (d_k + c_k)

とまとまります。i の子ノードが 3 つ以上ある場合は、任意の 2 個の組み合わせについて、すべて合計すればOKです。

ここまでをまとめると、次の実装になります。

import sys
from collections import defaultdict, deque

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

# dp[n] = (c, d) : n を頂点とするサブツリーについての (ノード数, 「子孫 i から n の距離」の合計)
  dp = [(0, 0)] * (N+1)
  dists = 0

  q = deque()
  q.append((1, 0)) # (node, parent)
  while q:
    i, p = q.pop()
    if i < 0:
      c, d = 0, 0
      for j in link[-i]:
        if j == p:
          continue
        c_j, d_j = dp[j]
        c += c_j
        d += d_j + 1*(c_j) # j->i の距離 1 を各子孫について加える
      dp[-i] = (c+1, d)

      links = []
      for j in link[-i]:
        if j == p:
          continue
        links.append(j)

      for index1 in range(len(links)):
        j = links[index1]
        c_j, d_j = dp[j]
        for index2 in range(index1+1, len(links)):
          k = links[index2]
          c_k, d_k = dp[k]
          dists += (d_j+c_j)*c_k + c_j*(d_k+c_k) # k -> i -> j の距離
      continue

    q.append((-i, p)) # 帰りがけ順で処理するための番兵

    for j in link[i]:
      if j == p:
        continue
      q.append((j, i))

  for n in range(1, N+1):
    dists += dp[n][1] # k -> i の距離をすべて加える
  print(dists)

main(sys.stdin)

もう一歩だけ最適化

ただし、上記のコードを提出すると、一部のケースで TLE になります。

i の子ノードが大量にある場合、「任意の2つの組み合わせ」の場合の数が大きくなるためです。ただ、実際の計算内容をよく見ると、一方の子ノード j を固定すると、もう一方の子ノード k からの寄与は、c_k および d_k それぞれの合計としてまとめることができます。そこで、これらの部分和を事前に計算しておけば、i ごとに個別に k のループを回す必要がなくなります。これを実装したのが、最終的なこちらの解答になります。

atcoder.jp