🍫

ABC 445 D 問題：2 つの順序を同時に扱うためのインデックスソートとポインタ

に公開
Python
競技プログラミング
アルゴリズム
idea

はじめに

ABC445 D（長方形の復元問題）について、復習を兼ねて本問から学んだことを共有します。対象は茶～緑レベルの想定です。

解法の大まかな方針として「復元する長方形の高さまたは幅と一致するピースを順に配置する」ことは比較的早く考えつきました。
しかし、毎回すべてのピースを走査して条件に合うものを探すと、計算量が \mathcal{O}(N^2) となって TLE してしまいます。
高速化するには高さ順と幅順の 2 種類の順序を同時に扱う必要がありますが、当初の私はその実装方法がわからずつまずいていました。

本記事では公式解説で用いられている「データ本体は並び替えず、インデックスのリストをソートし、ポインタで未使用要素を管理する」手法について、できるだけ丁寧に説明します。

問題

問題全文：D - Reconstruct Chocolate

サイズ H \times W の 1 つの長方形を分割して作られた N 個の長方形のピースが与えられるので、これらのピースを配置してもとのサイズの長方形に復元する問題です。

方針の整理

いま復元しようとしている長方形のサイズを (h, w) とします。初期状態では (h, w) = (H, W) です。
与えられる入力はこの長方形を縦または横に切って分割した結果なので、各ステップにおいて「高さが h と等しいピース」「幅が w と等しいピース」が必ず 1 つ以上存在します。
したがって、次の操作をピースがなくなるまで繰り返せばよいです。

  • 残っているピースから、高さが h と等しい、または幅が w と等しいものを 1 つ選ぶ
  • 選んだピースを原点の対角側の端に配置する
  • 復元する長方形のサイズ (h, w) をそのピースの分だけ狭める

全探索による実装

以上の方針を全探索で実装すると、次のような処理になります。

  • 以下の処理を N 回繰り返す
    • 残っているピースについて先頭から走査する
    • 高さが h と等しい、または幅が w と等しいピースを見つけたらそれを配置し、使用済みとする

この実装の計算量を求めます。
最初に N 個のピースがあり、1 個配置するごとに未使用のピースが 1 つ減ります。
したがって、ピースを探すための判定回数は少なくとも

N + (N - 1) + (N - 2) + \dots + 1 = \cfrac{N(N + 1)}{2}

となります。よって計算量は \mathcal{O}(N^2) で、たとえ使用済みのピースを探索範囲から除いたとしても問題の制約下で TLE になります。

ピースをサイズで降順ソートすることを考える

全探索を避けるために、たとえばピースを高さの降順に並べ替えることを考えます。
こうすれば「高さが h と一致するピース」を高速に見つけられそうです。

しかし、実際には「高さが h と一致するピース」を配置した次に「幅が w と一致するピース」を配置すべき場合があります。
ピースを高さ順に並び替えただけでは「幅が w と一致するピース」を高速に取り出すことができません。

つまり、必要なのは「高さの降順」「幅の降順」という 2 種類の順序を同時に持つことです。
これを実現するために インデックスのリスト を使う方法があります。

インデックスのリストで 2 つの順序を同時に管理する

インデックスのリストの実体は (0, 1, \dots, N - 1) という整数の並びです。これを元のリストの添字として使います。

本問ではピースのリスト pieces に対して「高さの降順」「幅の降順」という 2 種類の順序を同時に持ちたいのでした。
そのため、それぞれに対応するインデックスのリスト ord_h と ord_w を用意し、pieces の高さと幅をキーとして ord_h および ord_w をソートします。
ここで重要なのは pieces 自体を並び替えないということです。

# 配置するピースのリスト pieces: list[tuple[int, int]]
# pieces のサイズは N と等しい

ord_h = list(range(N))
ord_h.sort(key=lambda i: pieces[i][0], reverse=True)

ord_w = list(range(N))
ord_w.sort(key=lambda i: pieces[i][1], reverse=True)

たとえば pieces[ord_h[0]] は高さが最大のピースを、pieces[ord_w[0]] は幅が最大のピースを表します。

ポインタを利用して未使用の要素を管理する

次に、使用済みのピースを再利用しないように、ord_h, ord_w 上で未使用のピースを指すポインタ ih, iw を用意します。
ここでいうポインタとは、インデックスのリストで「いま見ている位置」を表す変数のことです。

ans = [(-1, -1)] * N
used = [False] * N
ih = 0
iw = 0

for _ in range(N):
    # ih, iw が未使用のピースを指すまでインクリメントする
    while ih < N and used[ord_h[ih]]:
        ih += 1
    while iw < N and used[ord_w[iw]]:
        iw += 1

    # ピースを配置し (h, w) を縮める
    if pieces[ord_h[ih]][0] == h:
        i = ord_h[ih]
        ans[i] = (1, w - pieces[i][1] + 1)
        w -= pieces[i][1]
        used[i] = True
    else:
        # pieces[ord_w[iw]][1] == w が成り立つ
        i = ord_w[iw]
        ans[i] = (h - pieces[i][0] + 1, 1)
        h -= pieces[i][0]
        used[i] = True

インデックスソートを用いた場合の計算量

ih, iw は while ループによって増加するのみで、減少することはありません。それぞれ最大でも N 回しか増加しないため、全体で while ループが処理される回数は高々 2N 回です。

ピースの配置は N 回ですが、インデックスのリストをソートするために計算量 \mathcal{O}(N \log N) がかかるため、全体の計算量は \mathcal{O}(N \log N) となります。これは十分高速です。

まとめ

本記事で解説したインデックスソートは次のような場面で有用です。

  • 複数の基準でソートした状態が必要である
  • 後で元データのインデックスが必要である

本問は「高さと幅の 2 種類の順序を同時に扱う必要がある」「出力を入力に合わせる必要がある」と、2 つの条件を満たしています。

ソート条件が複雑であったり、元のデータ順を並び替えたくないような場合は「代わりにインデックスのリストをソートする」という今回取り上げた手法が取れるかもしれません。
また、複数のインデックスのリストを先頭から参照していくためには、使用済みフラグとポインタを組み合わせて使うのが有効です。

参考

Discussion