AtCoder ABC241 B - Pasta 最大流やmapや総当たりで解いたり、unordered_mapは気を付けないと検証したり

はじめに

この問題をいろいろな解き方で解きます。

• 解法1: ブルートフォース O(NM)
• 解法2: Python: defaultdict O(N+M)
• 解法3: map C++ O(MlogN)
• 解法4: 注意: unordered_map C++ O(N+M)だが最悪O(MN)
• 解法5: sortされているならばO(N+M)で解く (ソートされている入力だった場合に有用)
• 解法6: 最大流で解く (二部グラフの最大マッチング1)
• 解法7: 最大流で解く2

■解法1,2,3は思い浮かびやすい解法です。
■解法4は解法3に似ていますが、ハッシュ衝突の危険性があります。
■解法5はLeetCodeっぽい解き方です。今回はO(NlogN + MlogM)なので嬉しくないです。
■解法6, 7は今回の問題に対してはtoo muchですし計算量も良くないですが、こういう問題で練習しておくと条件が複雑になったとき（今はBはAがx(というか同じ)時に条件を満たすが、BはAがxあるいはyでいいよという条件になったりしたとき）にも対応できて良いかもしれません。

解法1: ブルートフォース O(NM)

リストのままAに各Bが存在するかを調べ、存在したら消します。探索にO(N), 削除にO(N)かかるのでO(MN)になります。

実装(Python)

import sys
n, m = map(int, input().split())
data = list(map(int, input().split()))
datb = list(map(int, input().split()))
for x in datb:
    if data.count(x) == 0:
        print("No")
        sys.exit(0)
    del data[data.index(x)]
print("Yes")

解法2: Python: defaultdict O(N+M)

Aが何個あるかを数えておき、Bを調べるごとに-1します。足りなくなったら失敗です。O(N+M)で動作します。
Pythonの場合はハッシュ衝突が発生しにくく、のちに述べる解法4のような問題は起きません。

実装(Python)

import sys
n, m = map(int, input().split())
data = list(map(int, input().split()))
datb = list(map(int, input().split()))
from collections import defaultdict
d = defaultdict(int)
for x in data: d[x] += 1
for x in datb:
    if d[x] <= 0:
        print("No")
        sys.exit(0)
    d[x] -= 1
print("Yes")

解法3: map C++ O((N+M) logN)

計算量は@itchy_kneeさんからご指摘いただきました。ありがとうございます。また、$M \leq N$という制約があるので、O(NlogN)とも書けるともコメントいただきました。ありがとうございます！

解法2をC++で解きます。default_dictの代わりにmapを使いますが、操作がO(1)でなくO(logN)になるので時間計算量はO((N+M)logN)になります。
Aの挿入にNlogN, Bの検索にMlogNです。

実装(C++)

int main() {
    map<int, int> mp;
    int n, m; cin >> n >> m;
    int x;
    REP(i, n) { cin >> x; mp[x]++; }
    REP(i, m) {
        cin >> x;
        if(mp[x] == 0) { cout << "No" << "\n"; return 0; }
        mp[x]--;
    }
    cout << "Yes" << "\n";
    return 0;
}

解法4: 注意: unordered_map C++ 最悪O(MN)

以下は今回の問題の制約では問題(TLE)になりません。もしも、という考察です。

さて、解法3を見て、mapの計算量が気になるならunordered_map使えば各操作O(1)だよ！...というのは一見正しく聞こえます。この問題の制約では問題ないのですが、以下のような点に気を付けてください。
もしも、$n, m \leq 2 \times 10^5$などだとハッシュ衝突でTLEすることがあります。(Atcoderでそのようなテストケースはないかもしれませんが)。詳しくは以下を見てください。

試していきます。以下のコードを書いたとしましょう。(解法3のmapをunordered_mapにしただけです)

これを、AtCoderのカスタムテスト C++(Clangではない)で適当なn,m=200000のケースにかけると、以下のようになりました。(適当な乱数を生成しました)

実行時間 26 ms

ここで、以下のコードで生成した入力を入れてカスタムコードを実行します。

入力の生成(Python)

k=30727
print(20000, 20000)
a = [k*i for i in range(20000)]
b = a[::-1]
print(*a)
print(*b)

すると以下のようになりました。

実行時間 3138 ms

mapの際に、30727の倍数はハッシュを衝突させられる値です(これはmapの長さやコンパイラの条件によって変動するのでこの数で必ず発生するわけではないです)。今回は$n,m \leq 1e3$なので問題はないものの、CodeForcesやTopCoderのようにHackフェイズのあるコンテストでunordered_mapやunordered_setを使う際は気を付けましょう。(なお、対策は上記の記事上に書いてあります)

解法5: sortされているならばO(N+M)で解く (ソートされている入力だった場合に有用)

この問題が、「ソートされていた」時の解き方を考えます。今回はソートされていません。

• A, Bはソートされています。(今回はソートします)
• AとBのindex0にそれぞれポインタi, jを置きます。
• BをインクリメントしながらAを見ていき、Aが小さければi++します。Aが大きければNoです。Aが同じなら++します。Aがインクリメントできないとき、Noです。
• Bが最後までたどり着ければYesです。
  これは、入力がソートされている場合、時間計算量O(N+M)で空間計算量O(1)で解くことができます。LeetCodeなどでは頻出です。今回はわざわざsortしないといけないのでO(NlogN + MlogM)です。

実装(Python)

n, m = map(int, input().split())
data = list(map(int, input().split()))
datb = list(map(int, input().split()))
data.sort()
datb.sort()
i = j = 0
while i < n and j < m:
    if data[i] == datb[j]:
        i, j = i+1, j+1
        continue
    if data[i] > datb[j]: break
    i += 1
if j == m: print("Yes")
else: print("No")

解法6: 最大流で解く (二部グラフの最大マッチング1)

この問題は次のような二部グラフの最大マッチング問題と言い換えられます。

グループAのn人とグループBのm人がおり、全員数字の書いてある紙を持っています。
お互いに、同じ数字を持っている人とペアになりたいです。
m組のペアが作れる場合はYes, 作れない場合はNoを答えてください。

次のように考えます。

• 座標圧縮は行わなくてよいです。
• 先に、Bのノードの数字を探索しておき、g<int, vector<int>>で、[持っている数字] = [それを持っているノードのリスト]を持ちます
• Aを走査し、対応する数字のB側のノードに対してcap1の辺を張ります。(これはINFでもよいです)
• sからAの各ノードにcap1の辺を張り、Bの各ノードからtにcap1の辺を張ります。
• s,tにINF流し、tにm届いていればYesです。

実装(C++ w/ACL)

int main(int argc, char *argv[]) {
    int n, m; cin >> n >> m;
    int nodes = n + m;
    int nodet = n + m + 1;
    int x;
    map<int, vector<int>> g;
    mf_graph<int> mf(n + m + 2);
    REP(i, n){
        cin >> x; g[x].push_back(i);
        mf.add_edge(nodes, i, 1);
    }
    REP(i, m){
        cin >> x;
        for(auto &a: g[x]) mf.add_edge(a, n+i, 1);
        mf.add_edge(n + i, nodet, 1);
    }
    int matchCnt = mf.flow(nodes, nodet);
    if(matchCnt == m) cout << "Yes" << "\n";
    else cout << "No" << "\n";
}

解法7: 最大流で解く2

解法5は最大で辺がNM張られてグラフが大きくなります。例えば、すべての数字が同じ時です。そのため、以下のようにも作ってます。

• A, Bを先読みして、数字ごとにノード番号を採番します。座標圧縮のようですが、ソートしなくてよいです。
• 各ノードは何個数字を持つかを数えます
• sからA側のノードにcap =そのノードで辺を張ります。B側のノードが存在するなら、それに対してcap = そのノードで辺を張ります。
• B側のノードからtにcap =そのノードで辺を張ります。
• s,tにINF流し、tにm届いていればYesです。

実装(C++ w/ACL)

int main(int argc, char *argv[]) {
    int n, m; cin >> n >> m;
    vector<int> a(n);
    vector<int> b(m);
    map<int, pair<int, int>> aToNode;
    map<int, pair<int, int>> bToNode;
    int x;
    int nodecnt = 0; // 今から使っていいノード番号
    REP(i, n) cin >> a.at(i);
    REP(i, m) cin >> b.at(i);
    for(auto &x: a){
        // A側で定義されていないノードならそれを作って、[ノード番号, 0]で初期化
        if(aToNode.count(x) == 0){
            aToNode[x] = {nodecnt, 0};
            ++nodecnt;
        }
        aToNode[x].second++; // そのノードの数を++
    }
    for(auto &x: b){ // Bについてもおなじことをする
        if(bToNode.count(x) == 0){
            bToNode[x] = {nodecnt, 0};
            ++nodecnt;
        }
        bToNode[x].second++;
    }
    atcoder::mf_graph<int> mf(nodecnt+2); // nodecnt + 2個(s,t)が必要
    int nodes = nodecnt; // start
    int nodet = nodecnt + 1; // goal
    for(auto &x: aToNode){ // aの各ノードの処理
        mf.add_edge(nodes, x.second.first, x.second.second); // s -> aのノードにcap=その数の辺
        if(bToNode.count(x.first) > 0){
            // もし、対応するBのノードがいれば、辺を作る。capは同じ。
            mf.add_edge(x.second.first, bToNode[x.first].first, x.second.second);
        }
    }
    for(auto &x: bToNode) {
        // Bのノードからtに張る。capはB側のそのノードの数が限界
        mf.add_edge(x.second.first, nodet, x.second.second);
    }
    int matchCnt = mf.flow(nodes, nodet);
    if(matchCnt == m) cout << "Yes" << "\n";
    else cout << "No" << "\n";
}