AtCoder ABC245 E - Wrapping Chocolate の考え方と2通りの解き方: multiset, Fenwick Tree + 座標圧縮で解く

考え方を述べた後、C++のmultisetと、(multisetのない)Pythonの座圧+BITの解き方を示します。

基本的な方針0: 小問題1

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

• 全て同じ高さの長さwのチョコがn個あります。そして長さwの箱がm個あります。
• 箱がチョコの長さ以上(inclusive)の時、チョコは箱に入ります。

今、チョコが、[3,2,5]で箱が[1,2,10,4]だとします。この時、以下のように解けます。

• チョコを順にみていき、a:入る一番小さな箱を使います。

ここで、aをどう調べればよいでしょう。最も単純な考え方は残っている箱をすべてforで見ていけば$O(nm)$ですが、これは工夫することが可能です。

• 両者を昇順にソートします。[2,3,5]と[1,2,4,10]になります。
• 両方の先頭にポインタを持ちます。チョコポインタを順番に見ていき、それ以上の値を持つ箱が見つかるまで+1していきます。

このようなアプローチにすると、ソートがネックになり、$O(NlogN + MlogM)$で計算できます。

...という説明をしておいてなんですが、この問題ではこのアプローチは適応できません。

基本的な方針1: 小問題2

小問題1では最初にすべてのチョコの高さが分かっていました。これが順番に与えられる、つまり、

• n個のチョコがあることはあらかじめわかっているが、それぞれのチョコの長さは1クエリずつ飛んでくる。
• そして、毎回使う箱を答えたい。(一度使った箱は使えない)

というinteractiveな問題だとします。例えば、

• [1,2,10,4]の箱があり、3つのチョコがくると分かっている
• まず、3のチョコがきた。4か10を返せばいい
• 次に、2のチョコがきた。2か4か10を返せばいい(が、それまでに使ったチョコは使えない)

とします。この場合、最も単純な考え方は毎回O(M)でチョコの長さ以上の最も小さな箱を求めてO(NM)としても良いですが、以下のように工夫できます。

• 最初にMをソートしておく。二分探索を用いてチョコ以上の最も小さな箱を見つける。その箱を消す(=使ったものとする)。

これを行えるデータ構造としてはmultisetがあり、また、Fenwick Treeを使うことも可能です。このようにすれば、O(MlogM + NlogM)で解けます。

基本的な方針2:この問題の場合

しかし、この問題ではこれまでの問題では長さだけを考えていたものと違い高さがあります。そこで、以下のように処理します。

• チョコレート(choco)も箱(box)もh, wのpairなどで持ち、降順にソートします
• 各チョコについて以下の処理をします
  • 箱の先頭から見ていき、自分のh以上のhを持つ箱を使える候補とし、wを候補リストとして管理します
  • 今、候補リストにある箱はすべてhについてはチョコにとってOKな候補です(hについて降順でソートしているから)
  • このため、チョコのw以下の一番小さな箱を使います

この狙いは以下の通りです。

• ベースの考え方は小問題2と一緒です。その時に使える箱の中で一番小さなものを選んでいきます。
• ただし、徐々に使える箱も増えていきます(大きなhのチョコから見ていくので、進むにつれて使える箱は増えます。

multiset

手順通りに実装を行っていきます。

• pairの降順のソートは昇順にソートしてやり、backで参照、pop_backで消せばよいです。(照準を後ろからみることは降順と同値です)
• multiset candidateがそのチョコがいられれるhを持つwのmultisetです。lower_boundすることで「その値以上の最小の値」のイテレータを返します。
• multiset.erase()は以下２つの使い方があるので注意します。今回はイテレータのまま扱います。
  • 1: erase(値) はその値のエントリをすべて削除します
  • 2: erase(イテレータ) はその値のエントリだけ削除します

multisetは各計算量を以下で行えます。

• 追加: $O(logN)$
• 削除: $O(logN)$
• lower_bound: $O(logN)$

実装(C++)

int main() {
    FASTIOpre();
    int n,m; cin >> n >> m;
    vector<int> A(n), B(n), C(m), D(m);
    REP(i, n) cin >> A.at(i);
    REP(i, n) cin >> B.at(i);
    REP(i, m) cin >> C.at(i);
    REP(i, m) cin >> D.at(i);
    vector<pair<int, int>> choco, box;
    // h x wのpairをつくる　
    REP(i, n) choco.push_back({A.at(i), B.at(i)});
    REP(i, m) box.push_back({C.at(i), D.at(i)});
    // 昇順でsortする
    sort(choco.begin(), choco.end());
    sort(box.begin(), box.end());
    // hがいまのチョコより大きい箱のwのset
    multiset<ll> candidate;
    // チョコを後ろから見ていく(hが大きいものから見ていく)
    while(choco.size() > 0) {
        auto &[a, b] = choco.back();
        choco.pop_back();
        while(box.size() > 0){ /* 使えるはこの取り出し */
            auto &[c, d] = box.back();
            if(c >= a){ /* まだhが今のチョコより大きい場合 */
                candidate.emplace(d);
                box.pop_back();
            }
            else break;
        }
        /* このチョコを入れるのに十分な一番wの小さな箱のイテレータを得る */
        auto it = candidate.lower_bound(b);
        if(it == candidate.end()){
            cout << "No" << "\n";
            return 0;
        }
        candidate.erase(it); /* multiset.eraseはitを消すときは単一エントリを消すだけ */ }
    cout << "Yes" << "\n";
    return 0;
}

(シンプルな)Fenwick Tree + 座標圧縮

Fenwick Treeを使うことでもこの問題を解けます。Fenwick TreeはNこの要素の１次元配列(1-indexedで)要素1からkまでの和(など)を$O(logN)$で求めることができます。これを$sum(k)$とします。配列$a[x]$を長さ$x$の箱の数、として扱えば、
ある長さ$l$以上の箱が何個あるかは$sum(n) - sum(l-1)$で求めることができます。
つまり、二分探索をしてやり、これが0にならない最小のlを探すことで、使うべき箱がlとわかります。
これには、$O(log^2N)$の計算量がかかります。以下のようにFenwickTree上でlower_boundを実装で計算量は落とせます。

尚、この問題では与えられる長さが$10^9$と長いため座標圧縮します。

実装(Python)

# https://ikatakos.com/pot/programming_algorithm/data_structure/binary_indexed_tree
class Bit:
    def __init__(self, n):
        self.size = n
        self.tree = [0] * (n + 1)

    def sum(self, i): # 1-origin, [i,x]
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & -i
        return s

    def add(self, i, x):
        while i <= self.size:
            self.tree[i] += x
            i += i & -i

def do():
    n, m = map(int, input().split())
    data = list(map(int, input().split()))
    datb = list(map(int, input().split()))
    datc = list(map(int, input().split()))
    datd = list(map(int, input().split()))

    l = []
    for x in datb:l.append(x)
    for x in datd:l.append(x)
    # 座標圧縮する
    zatsu = sorted(set(l))
    zatsuTable = dict()
    zatsuTableRev = dict()
    for ind, value in enumerate(zatsu):
        zatsuTable[value] = ind
        zatsuTableRev[ind] = value
    newl = []

    choco = []
    box = []
    for i in range(n):
        a,b = data[i], datb[i]
        choco.append( (a,zatsuTable[b] +1) ) # bitするので+1
    for i in range(m):
        c, d = datc[i], datd[i]
        box.append( (c, zatsuTable[d] +1) ) # bitするので+1

    bit = Bit(len(zatsuTable) + 1)
    bitmax = len(zatsuTable) + 1

    降順ソート
    choco.sort(reverse=True)
    box.sort(reverse=True)

    from collections import deque
    choco = deque(choco)
    box = deque(box)
    while len(choco) > 0:
        a, b = choco.popleft()
        #print("try ", a,zatsuTableRev[b-1])
        # 使えるboxを出してくる
        while len(box) > 0:
            c, d = box[0]
            if c >= a: # 使える
                c, d = box.popleft()
                #print("can use box", c, zatsuTableRev[d-1])
                bit.add(d, 1)
            else: # 使えない場合
                break # 取り出さずにbreak
        canuse = bit.sum(bitmax) - bit.sum(b - 1) # bよりも大きいものが1つでもあるか？
        if canuse == 0: # もしないならダメ
            print("No")
            return

        # (ng, ok] for min value
        def func(mid):
            canuse = bit.sum(mid)
            canuse -= bit.sum(b - 1)
            if canuse > 0: return True
            return False
        ok = bitmax
        ng = b - 1
        while (abs(ok - ng) > 1):
            mid = (ok + ng) // 2;
            if(func(mid)) :ok = mid;
            else : ng = mid;
        #print("use", ok, bitmax, b-1)
        bit.add(ok, -1)
    print("Yes")
do()