二部マッチングでTLEに苦しんだ話

opencupで二部マッチングを使う問題があって、自分のライブラリが遅くてTLEにハマりました。
そのときに得た知見をまとめておきます。
頂点は合計200,000個、辺は200,000本、TLは3秒で、writerや他の人は400msくらいで通ってたらしい。

続きを読む

木と計算量 後編 〜全方位木DP〜

前編はこちら

前編は汎用性も高い話題でしたが、後半は少しマニアックで雑学的です。

予習

全方位木DPについて予習します。

補足として、以下の問題を解く具体的な実装も置いておきます。

N 頂点の木が与えられる。
頂点 v を含む部分木の個数を全ての v について求めよ。
制約:1 ≦ N ≦ 10^6

こういう、根を全部試して根付き木の問題を一括で計算する、的な状況で全方位木DPが活躍します。

pythonで実装するとこんな感じになります。

def dfs(v, p=-1):
  deg[v] = len(to[v])
  res = 1
  dp[v] = [0]*deg[v]
  for i in range(deg[v]):
    u = to[v][i]
    if u == p:
      pi[v] = i
      continue
    dp[v][i] = dfs(u, v)
    res *= dp[v][i]+1
    res %= mod
  return res

def bfs(v, res_p=0, p=-1):
  if p != -1:
    dp[v][pi[v]] = res_p
  dpl = [1]*(deg[v]+1)
  for i in range(deg[v]):
    dpl[i+1] = dpl[i]*(dp[v][i]+1)%mod
  dpr = [1]*(deg[v]+1)
  for i in range(deg[v]-1,-1,-1):
    dpr[i] = dpr[i+1]*(dp[v][i]+1)%mod
  ans[v] = dpr[0]
  for i in range(deg[v]):
    u = to[v][i]
    if u == p:
      continue
    bfs(u, dpl[i]*dpr[i+1]%mod, v)

dfs(0)
bfs(0)

本編

本編と言いつつも、こちらは雑談みたいなものです。

全方位木DPには色んな流派があるみたいですね。
その中で、zerokugiさんの方法が計算量の視点から興味深かったので紹介します。

追記

僕のオススメは上に書いた方法です。(定数倍、実装量、ロジックの簡潔さなどの点から)
あと、zerokugiさんの方法っぽくやる場合でも次数の降順より行きがけ順の方が定数倍が良いです。

さっきの問題をzerokugi法で実装すると以下のようになります。

def dfs(v, p=-1):
  res = 1
  for i in range(deg[v]):
    u = to[v][i]
    if u == p:
      pi[v] = i
      continue
    if dp[v][i] is None:
      dp[v][i] = dfs(u, v)
    res *= dp[v][i]+1
    res %= mod
  return res

roots = []
for v in range(N):
  deg[v] = len(to[v])
  dp[v] = [None]*deg[v]
  roots.append([deg[v], v])
roots.sort()
roots.reverse()
for _, v in roots:
  ans[v] = dfs(v)
  dpl = [1]*(deg[v]+1)
  for i in range(deg[v]):
    dpl[i+1] = dpl[i]*(dp[v][i]+1)%mod
  dpr = [1]*(deg[v]+1)
  for i in range(deg[v]-1,-1,-1):
    dpr[i] = dpr[i+1]*(dp[v][i]+1)%mod
  for i in range(deg[v]):
    u = to[v][i]
    dp[u][pi[v]] = dpl[i]*dpr[i+1]%mod
print(ans)

実はこのアルゴリズムの計算量は O(N) なのです。

丁寧に書くと長くなるので、以下に示す本質部分だけを解析します。

  • dfs内のfor i in range(deg[v]):

dfs(v)を呼ぶ度にO(deg[v])がかかるという訳です。

根として呼ばれるのは各 v について1回ずつなので O(N) です。(Σ(deg[v]) =辺の本数*2)
問題はdfs内から呼ばれる場合です。

dfs(v)がdfs(u)内から呼ばれるのは、各辺 u-v について高々1回です。(結果がメモされるため)
で、実際に呼ばれる条件は以下の通りです。

  • u 側の部分木内に次数が deg[v] 以上の頂点がある

別の言葉で言うと「v より先に根としてdfsを呼ばれる頂点が u 側の部分木内に存在する」です。

つまり、以下のことが言えれば計算量が O(N) であることが示せます。

頂点数 N のどんな木についても、Σ(「vを取り除いてできる部分木のうち、次数がdeg[v]以上の頂点が含まれるものの個数」* deg[v]) が O(N) である。

適当な頂点を根とした根付き木だとみなします。上記の式は以下のように言い換えられます。

Σ(「vの下に付く部分木のうち、次数がdeg[v]以上の頂点が含まれるものの個数」*「vの子の個数」)

実際には根方向の辺も考えないといけないので「vの下〜個数」には +1 が付くこともありますし、「vの子の個数」にも +1 は付きますが、Σ(deg[v]) が O(N) であることを思い出すと、この部分につく定数は無視できることが言えます。

さて、これの最大値の抑え方が分からなかったのでtwitterで助けを求めると、hosさんからエレガントな解答をいただいたので紹介します。(WA_TLEさんもありがとうございます)

頂点数 N 最大次数 K の根付き木についての最大値を f(N,K) とおくと、f(N,K) ≦ N-1-K。

これで、下に示すような帰納法が回るようになります。("次数"に親方向の辺は数えないものとします)
自分も帰納法が回せたら良いなぁとは思っていましたが、この式には辿り着けませんでした。すごい。

とりあえず、f(1,0) = 0
根の次数を d とすると、各 d に対して以下が成り立つ。
f(N,K) \leq max(\sum (f(n_i,k_i)+min(k_i,d))\ | \sum n_i = N-1,\ k_i \lt n_i,\ k_i \leq K)
右辺のmaxの中身の f を外すと、
\sum n_i - d - \sum k_i + \sum min(k_i,d) = (N - 1) - d - \sum max(k_i-d,0)
d = K の場合:そのまま N - 1 - K 以下だと言える。
d < K の場合:k_i のいずれかは K なため d + \sum max(k_i-d,0) \geq K であり、やはり N - 1 - K 以下だと言える。

min(k_i,d) の部分が少しテクニカルですが、k_i と d の大小で場合分けしてみると正しいことが確認できるかと思います。

あとがき

次数の降順でやれば計算量が O(N) で抑えられるってかなり奇妙で面白い!

木は奥深いですね。

おまけ

全方位木DPについて、普通のDPの場合はこれで良いのですが、データ構造で高速化するDPの場合などはまた別のアプローチをする必要があったりします。(前にHackerRankで使って以来、使う問題とは出会ってないけど。これの記事需要あります?)

木と計算量 前編 〜O(N^2)とO(NK)の木DP〜

この記事はCompetitive Programming Advent Calendar 2018の46日目の記事として書かれました(嘘)

最近、木上のアルゴリズムの面白い計算量解析が2つ話題になったのでまとめておきます。

予備知識

まず、https://web.archive.org/web/20150819082918/https://topcoder.g.hatena.ne.jp/iwiwi/20120428/1335635594 について復習します。
iwiさんのブログとは違う、より直感的な解析方法も紹介します

以下の問題を考えます。

N 頂点の木が与えられる。
頂点 1 を含む頂点数 K の根付き木の個数を求めよ。
制約:1 ≦ K ≦ N ≦ 3000

典型的な木DPの問題です。
解法は以下の通りです。(解法の細かい説明は本題ではないので追わなくて大丈夫です)

頂点 1 を根とした根付き木にして以下のようなdpテーブルをボトムアップに計算していく。
・dp[v][i] = 頂点 v を根とする頂点数 i の根付き木の個数

例えばpythonで書くと以下の通りです。

def dfs(v):
  sz[v] = 1
  dp[v] = [0]*(sz[v]+1)
  dp[v][1] = 1
  for u in to[v]:
    dfs(u)
    merged = [0]*(sz[v]+sz[u]+1)
    for i in range(sz[v]+1):
      for j in range(sz[u]+1):
        merged[i+j] += dp[v][i]*dp[u][j]%mod
    sz[v] += sz[u]
    dp[v] = merged
  dp[v][0] = 1

(変数宣言、入出力などは省略しています)
このアルゴリズムの計算量について考えていきます。

最も重い部分は i,j に関する二重ループの部分です。
ここの計算量は O(sz[v] * sz[u]) です。
つまり、サイズがそれぞれ a, b のdpテーブルをマージするときに O(ab) の計算量がかかっています。
これらを合計すると一見 O(N^3) の計算量が掛かるように思えます。
しかし、実は全体で O(N^2) になっているのです。

この計算量を解析するために以下のような問題を考えます。

N 個のグループがあり、初めは各グループに頂点が1つずつ含まれています。
これらをマージしていき、最終的に1つのグループにしたいです。
グループ A,B をマージするとき、A に含まれる頂点と B に含まれる頂点の間を結ぶような辺を全て追加します。(つまり、|A|*|B| 本の辺を追加します)
マージの順番を工夫したとき、追加する辺の本数は最大で何本でしょうか?

例えば以下のような流れになります。

答えは「どのような順番でマージしても完全グラフになるので、N*(N-1)/2」でした。
(どの2頂点間の辺についてもちょうど1回ずつ追加されるため)

この問題が計算量解析にどう関係しているかは、以下の2つを比較すれば分かるでしょう。

  • サイズがそれぞれ a, b のdpテーブルをマージするときに O(ab) の計算量がかかる
  • サイズがそれぞれ a, b のグループをマージするときに ab 本の辺を追加する

辺の本数がそのまま計算量を表しているのです。
というわけで先ほどのアルゴリズムの計算量は O(N^2) なのでした。

本編

予備知識だけでもそれなりのボリュームでしたが続けます。

以下のような問題を考えます。

N 頂点の木が与えられる。
頂点 1 を含む頂点数 K の根付き木の個数を求めよ。
制約:1 ≦ N ≦ 10^5, 1 ≦ K ≦ 500

先ほどと同じようなDPをpythonで書くと以下のようになります。

def dfs(v):
  global ans
  sz[v] = 1
  dp[v] = [0]*(sz[v]+1)
  dp[v][1] = 1
  for u in to[v]:
    dfs(u)
    merged = [0]*(sz[v]+sz[u]+1)
    for i in range(sz[v]+1):
      for j in range(sz[u]+1):
        merged[i+j] += dp[v][i]*dp[u][j]%mod
    sz[v] += sz[u]
    dp[v] = merged
    if sz[v] > K:
      sz[v] = K
      dp[v] = dp[v][:K+1]
  if sz[v] >= K:
    ans += dp[v][K]
    ans %= mod
  dp[v][0] = 1

要点は、dpテーブルのサイズが K を超えたら K になるようにカットしている点です。
このアルゴリズムの計算量を解析するために、簡略化した以下の問題を考えます。

N 個の集合があり、初めは各集合のサイズが1です。
これらをマージしていき、最終的に1つの集合にしたいです。
グループ A,B をマージするとき、min(|A|, K) * min(|B|, K) のコストがかかります。
マージの順番を工夫したとき、コストの合計は最大でいくらになるでしょうか?

追記:下の方により簡潔な解析方法を書きました。

サイズが K 未満の集合を「小」、K 以上の集合を「大」と表すことにして場合分けをします。

小と大のマージ

小と大のマージコストは「(小のサイズ) * K」です。
(集合の要素)に注目すると、各元が小-大マージの小の元として選ばれる回数は高々1回です
また、元の個数は N なので Σ(小のサイズ) は N で抑えられます。
よって、小と大のマージコストの合計が O(NK) であることが言えました。

大と大のマージ

大と大のマージコストは K^2 です。
大と大を x 回マージするためには、大の集合が x+1 個必要です。
大の集合は高々 N/K 個しか作れないので、大と大のマージは高々 N/K 回です。
K^2 * (N/K)\ =\ NK なので、大と大のマージコストの合計が O(NK) であることが言えました。

小と小のマージ

小と小のマージだけを行うことを考えます。
各集合のサイズは 2K 以上にはなりえません。
また、各集合での合計コストは O(集合サイズの二乗) です。(予備知識参照)

つまり、小と小のマージコストの合計の最大値は 以下の式で上から抑えられます。

  • max(\sum A_i^2\ |\ A_i \leq K,\ \sum A_i = N)(A を正整数列とします)(定数倍は無視しています)

これは Ai が全て K の場合に最大となりK^2 * (N/K)\ =\ NK となります
これで、小と小のマージコストの合計が O(NK) であることが言えました。


これらをまとめると、この問題の答えは O(NK) となります。
これで、上記のアルゴリズムの計算量は O(NK) であることが言えました。

追記

より直感的で場合分けもない解析方法を知ったので書いておきます。
予備知識のところで使った手法に似ています。

マージの過程を二分木のような形で表します。

で、マージするときに左のグループ内の右からK個と右のグループ内の左からK個の頂点の間を張ることを考えます。

  • 各頂点間には高々1回しか辺が張られない
  • 距離が2K以上離れた頂点との間には辺は張られない

ということが言えるので、辺の本数の合計はO(NK)となります。

あとがき

結構長く競プロをやってるつもりだけど、この計算量今まで知らなかった。面白い。
場合分けをしない良い感じの解析方法を見つけたりしたら教えてください。(上の追記参照)

うむ,やっぱあまりに汎用性が高いので僕が知らなかっただけで常識なのではないかと思い始めた.

ちなみにこれを知ったきっかけはHello 2019 Gでした。
(※ そのままdpしてもダメなので式変形などで工夫をする必要はあります)

後編はこちら

最大マッチングを貪欲する問題の証明典型テク

ARC092C - 2D Plane 2N Pointsの証明が話題になってたっぽいので問題を見てみたら、最大マッチング系貪欲の証明テクが詰まった問題だなぁと思ったので書いてみた。

問題概要

  • 平面上に赤い点と青い点がある。
  • 赤い点が青い点の左下である(つまりxもyも小さい)とき、それらはペアにできる
  • 何ペア作れるか

解法

青い点をxの昇順に見て、マッチングできる赤い点があればそのうちyが最大のものとマッチングさせていく貪欲。

証明テク1:マッチングできるときはしてもよい

つまり、あえてマッチングさせないことによって答えが増えることはない。

なぜなら、ある青い点 A と赤い点 X をマッチングできるとき、

  • マッチングさせる:残りの点集合から A と X が消える
    • メリット:答えが 1 増える
  • マッチングさせない:残りの点集合から A が消える
    • メリット:X を残せる

で、「X を残せる」ことによって答えは高々 1 しか増やせないため、(マッチングさせたときの答え) < (マッチングさせないときの答え) となることはない。

ただし、A と別の赤い点 Y をマッチングさせる方が答えが良くなるという可能性はある点に注意。

このテクは結構出てくる。
使えるのは以下のようなケースのとき。

  • 次数1の頂点がある:木でのマッチングなど
  • マッチング相手に全順序がある:後述

例題:最近だとこれとか

証明テクテク2:マッチング相手に最弱の頂点があれば、それとマッチングさせてよい

以下のような状況を例に取って説明します。

f:id:snuke:20190103234249p:plain

緑の線の接続関係は同じとします。(灰色は何でもいい)

今、A にマッチングできる点は X と Y です。
X, Y とマッチングさせられる頂点集合を n(X), n(Y) とすると、n(X) ⊆ n(Y) です。
このとき、A と X はマッチングさせてもいいということです。
(少なくとも X か Y のどちらかとマッチングさせてもいいというのはテク1から言えますね)

なぜなら、

  • Y とマッチングさせる:このとき残る辺集合を E とする
  • X とマッチングさせる:このとき残る辺集合は(X と Y を区別せずに見ると) E + 辺(Y,α) となる

となり、辺が多い方が当然答えは大きくなるので X とマッチングさせても良いため。

マッチング相手が3頂点以上あった場合でも、その中で n(v) が最弱の頂点があればそれとマッチングさせて良い(つまり、どの他の頂点 u についても n(v) ⊆ n(u) が成り立つ v があれば)
大抵の問題ではマッチング相手の中に全順序があるが、一応そこまでは必要ない。

相性の良い競プロテク1:平面走査で変数を1つ消す

この問題の最後のピースです。
このままだと2変数あるので全順序とかも成り立ってないので工夫が必要です。

"A" として試す青い点を x の昇順で見ていくことによって x を消して1変数にすることができます
すると y の大小だけの全順序が成り立つようになるわけです。

もう少し詳しく言うと、

  1. A の x 座標 x_A は残っている青い点のうち最小である
  2. 以降は x 座標が x_A 以下の赤い点は全部 x 座標が -∞ だと思っても答えは変わらない
  3. A とマッチングできる頂点は x 座標が -∞ の赤い点のうち y 座標が y_A 以下のものである
  4. x 座標が -∞ の赤い点の中では、y 座標は大きければ大きいほど "弱い"

ということから、最初に書いたような解法の正当性が示せる。

あとがき

貪欲法の問題は、正しいと確信できる程度の証明までして初めて解けた言えると思うんですが、
証明方法にも典型パターンというものは存在して、そういうのに積極的に触れていけばその内できるようになっていくと思います。