Zenn
🃏

ABC260 D - Draw Your Cards Python解答例

2022/07/18に公開
Python
競技プログラミング
AtCoder
競プロ
tech

AtCoder Beginner Contest 260 D - Draw Your CardsをPythonで解きます。

問題

問題文をAtCoderのページより引用します。

問題文

1からNが書かれたN枚のカードが裏向きで積まれた山札があり、上からi枚目のカードには整数P_iが書かれています。

この山札を使って、以下の行動をNターン繰り返します。

  • 山札の一番上のカードを引いて、そこに書かれた整数をXとする。
  • 場に見えている表向きのカードであって書かれた整数がX以上であるもののうち、書かれた整数が最小のものの上に、引いたカードを表向きで重ねる。もし場にそのようなカードがなければ、引いたカードをどれにも重ねずに表向きで場に置く。
  • その後、表向きのカードがK枚重ねられた山が場にあればその山のカードを全て食べる。食べられたカードは全て場から消滅する。

各カードについて、何ターン目に食べられるか、あるいは最後まで食べられないかを求めてください。

制約

  • 入力は全て整数
  • 1 \le K \le N \le 2 \times 10^5
  • Pは(1,2,\dots,N)の順列 ((1,2,\dots,N)を並べ替えて得られる列 ) である

解答例

場に出ているカードをうまく管理する

問題で指定されている処理を行うためには、以下のことについて管理できている必要があります。

  1. 場に見えている表向きのカード(以下、先頭のカードと呼びます)の集合
  2. 各先頭の下に重ねられているカードの集合(以下、グループと呼びます)の要素とサイズ

山札からカードをドローしたとき、1.の集合を使って「場に見えている表向きのカードであって書かれた整数がX以上であるもののうち、書かれた整数が最小のもの」(以下、このカードをYと呼びます)を探します。
また、2.を使って「表向きのカードがK枚重ねられた山」になったかどうかを判別し、カードを食べる処理を行います。

1.の集合は順序付き集合(C++でいうstd::set)か、あるいはセグメント木を使うことで効率的に管理することができます。 2.の集合はUnionFind木を使うことで高速に処理することができます。

SortedSetを使用した例

tatyam氏によるSortedSetを使用することにより、Pythonでも順序付き集合を扱うことができます。

先頭のカードに書かれている数値を順序付き集合の要素として管理することで、カードYを二分探索によって高速に求めることができます。
カードYが存在したらその上にXが書かれたカードを重ねますが、この操作は、

  • カードYを集合から取り除く
  • カードXを集合に加える

の一連の操作であるとみなせます。

また、2.の集合はUnionFind木を用いて管理します。カードYにカードXを重ねたとき、YとXをマージします。
カードXが属するグループのサイズはUnionFind木により管理することができます。これにより、グループのサイズがKを超えたかどうかを検知することができます。

UnionFind木は各集合の元を保持することはできませんが、各集合の根はわかります。この問題ではあるグループ内の元は全て同じタイミングで食べられるため、そのグループの根が食べられたタイミングを記録しておき、最後に全てのカードの数字iに対して、iが属するグループの根が食べられたターンを求めれば良いです。

SortedSetを用いた実装例を以下に示します。尚、各クラスの実装コードは長いので省略しています。

d.py

# SortedSetの実装(省略)

# UnionFind木の実装(省略)

def main():
    # 入力受け取り
    N, K = map(int, input().split())
    P: List[int] = list(map(int, input().split()))

    # 各カードが食べられたターンを記録しておくリスト(1-indexed)
    answer: List[int] = [-1] * (N + 1)

    # 順序付き集合
    S = SortedSet()
    # UnionFind木
    uf = UnionFind(N + 1)

    # 山札からカードを引いていく。iがターン数、pが引いたカードに書かれた数(X)。
    for i, p in enumerate(P):
        # 順序付き集合を二分探索してカードYのインデックスを探す
        index: int = S.index_right(p)
        # カードYが存在したとき:
        if index < len(S):
            # カードXとカードYを結合する
            uf.unite(p, S[index])
            # カードYは先頭のカードでなくなるので集合から削除する
            S.discard(S[index])
        # 順序付き集合にカードXを追加する
        S.add(p)

        # カードXを追加したことでグループのサイズがKを超えたら:
        if uf.get_size(p) == K:
            # 食べられたターンを、グループの根のカードにだけ記録する
            answer[uf.get_root(p)] = i + 1
            # 先頭のカードからカードXを削除する
            S.discard(p)

    # 全てのカードについて、それが食べられたターンは、
    # 根が食べられたタイミングと同じになる。
    for i in range(1, N + 1):
        answer[i] = answer[uf.get_root(i)]

    # 回答出力
    for a in answer[1:]:
        print(a)


if __name__ == "__main__":
    main()

実際の提出はこちら。

セグメント木を使った例

以下の図のような、インデックスと同じ値を要素として持つセグメント木を考えます。

要素の型は整数で、積は最小値を取る(min)操作とし、単位元は\inftyです。

カードXを集合に加えることは、対応するインデックスに当たる葉の値をXに置き換えることに相当します。
また、カードXを集合から削除することは、対応するインデックスに当たる葉の値を\inftyにすることに相当します。
また、カードYを探す操作は、区間[X, N]で積を取ることに相当します。

入力例1において、3ターン目の時点でのセグメント木は以下の図のようになります。

この時点で、場には2と5が書かれたカードがある状態です。
この中からカードYを探す操作は、赤枠で示した区間の中で最小値を取る操作に相当します。

3ターン目でカード3にカード2を重ねたとき、セグメント木は以下の図の状態になります。

グループのサイズが2になったので、カード2とカード3は食べられます。3ターン目終了時点で、セグメント木は以下の状態になります。

UnionFind木によるグループの管理は、SortedSetを使うときと同様です。

セグメント木を用いた実装例を以下に示します。クラス実装は省略しています。

d.py

# セグメント木の実装(省略)

# UnionFind木の実装(省略)

def main():
    # 入力受け取り
    N, K = map(int, input().split())
    P: List[int] = list(map(int, input().split()))

    # 単位元を返す関数
    def e() -> int:
        return 1 << 60

    # セグメント木の構築。初期状態では全ての葉は単位元が入る(集合が空である)
    seg = SegmentTree[int](lambda a, b: min(a, b), e, [1 << 60 for i in range(N + 1)])
    # UnionFind木の構築
    uf = UnionFind(N + 1)

    # 答えを格納するリスト(1-indexed)
    answer: List[int] = [-1] * (N + 1)

    # 山札からカードをドローしていく
    for i, p in enumerate(P):
        # カードYを探す
        t = seg.prod(p, N + 1)

        # カードYが存在したとき(= 積tが単位元でなかったとき)
        if t != e():
            # カードXとカードYを結合する
            uf.unite(p, t)
            # 集合からカードYを取り除く
            seg.set(t, e())
        # 集合にカードXを追加する
        seg.set(p, p)

        # カードXを追加したことでグループのサイズがKを超えたら:
        if uf.get_size(p) == K:
            # 食べられたターンを、グループの根のカードにだけ記録する
            answer[uf.get_root(p)] = i + 1
            # 集合からカードXを取り除く
            seg.set(p, e())


    # 全てのカードについて、それが食べられたターンは、
    # 根が食べられたタイミングと同じになる。
    for i in range(1, N + 1):
        answer[i] = answer[uf.get_root(i)]

    # 回答出力
    for a in answer[1:]:
        print(a)


if __name__ == "__main__":
    main()

実際の提出はこちら。

GitHubで編集を提案

Discussion