見出し画像
Photo by アンリ・ド・トゥールーズ＝ロートレック / メトロポリタン美術館

AtCoder精選良問「D - Draw Your Cards」Diff 1074

Qlitre

この記事はAtCoderの精選良問集マガジンの中の記事です。
私が良問だと思った問題を解説付きでまとめています。

問題リンク

問題概要

1,2,…Nの順列で構成された山札が与えられる。
そこから下記のクエリを処理していく。

  1. 山札の一番上のカードを引き、値をXとする。

  2. 場にあるカードの中で、X以上の値を持つ最小のカードを探す。

  3. もしX以上の値を持つカードが存在すれば、そのカードの上に引いたカードを表向きで重ねる。存在しなければ、引いたカードを表向きで場に置く。

  4. 場にあるカードがK枚以上になった場合、そのカードを全て食べる。

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

考え方

こういう問題はまずはシミュレーションを実際に行ってみる。
入力例である山札が[3, 5, 2, 1, 4]、K=2の場合。

  1. 山札から3を引く。場には何もないので、3を場に出す。

  2. 山札から5を引く。場に3が出ているが、5より小さいので、5も場に出す。

  3. 山札から2を引く。場に出ている3と5は2よりも大きい。最小の3に2を重ねる。2枚以上になったので、2と3を食べる。

  4. 山札から1を引く。場に残っている5が大きいので、5の上に1を重ねる。1と5を食べる。

  5. 山札から4を引く。場に残っているカードはないので、4を場に出す。4は食べられない

以上の流れになる。
データ構造で殴る系の問題。

まず一つのポイントが、山札を引いた時に、「自分より大きくて一番小さいカードを探す」という点。これは順序付き集合などのデータ構造を使えば実装できる。

次に、「カードを重ねる」という動作について。これも順序付き集合を用いて、重ねられる側のカードを集合から削除して、引いたカードを場に出せばいい（重ねられる側のカードが今後比較されることはないため）。

最後に場のカードが何枚重なっているか？という管理。
これは、レコードの機能をうまく使えば実装ができそう。

{ 一番上のカード番号: [含まれるカード1,含まれるカード2…] }

というように管理する。

例えば、上の例で3の上に2を重ねる時はこんなイメージ。

// 現在の場の状況
{3:[3],5:[5]}
// 2を引いたので3に重ねる
{3:[3], 5:[5], 2:[3,2]}

実装

TypeScriptでは順序付き集合は用意されていないので、自前で用意する必要がある。
とはいえ、そんな技量は自分にはない。
そこで、つよつよ競技プログラマーのtatyamさんがPython用に公開しているSortedSetクラスを、ChatGPTでTypeScriptに変換をかけて使わせてもらった。

import { createInterface } from "readline";
import * as fs from "fs";
import { log } from "console";
import { exit } from "process";

/**
 * 
 * @param arr ソート済みの配列
 * @param x 検索値
 * @returns 
 * 値 x が最初に現れるインデックスを返す。
 * x が arr 内に存在しない場合は、x を挿入するためのインデックスを返す
 */
export const bisectLeft = (arr: number[], x: number, lo: number = 0, hi: number = arr.length): number => {
    while (lo < hi) {
        const mid = Math.floor((lo + hi) / 2);
        if (arr[mid] < x) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    return lo;
}

/**
 * 
 * @param arr ソート済みの配列
 * @param x 検索値
 * @returns 
 * 値 x が最後に現れるインデックスを返す。
 * x が arr 内に存在しない場合は、x を挿入するためのインデックスを返す
 */
export const bisectRight = (arr: number[], x: number, lo: number = 0, hi: number = arr.length): number => {
    while (lo < hi) {
        const mid = Math.floor((lo + hi) / 2);
        if (arr[mid] <= x) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    return hi;
}


type T = any;

export class SortedSet {
    private static BUCKET_RATIO = 50;
    private static REBUILD_RATIO = 170;

    private size: number;
    private a: T[][];

    constructor(a: Iterable<T> = []) {
        let arr = Array.from(a);
        if (!arr.every((_, i) => i === arr.length - 1 || arr[i] < arr[i + 1])) {
            arr = Array.from(new Set(arr)).sort();
        }
        this.a = [];
        this.build(arr);
        this.size = arr.length;
    }

    private build(a: T[] = Array.from(this)) {
        const size = this.size = a.length;
        const bucketSize = Math.ceil(Math.sqrt(size / SortedSet.BUCKET_RATIO));
        this.a = Array.from({ length: bucketSize }, (_, i) => (
            a.slice(size * i / bucketSize, size * (i + 1) / bucketSize)
        ));
    }

    *[Symbol.iterator]() {
        for (const a of this.a) {
            for (const item of a) {
                yield item;
            }
        }
    }

    length() {
        return this.size;
    }

    private findBucket(x: T): T[] {
        for (const a of this.a) {
            if (x <= a[a.length - 1]) {
                return a;
            }
        }
        return this.a[this.a.length - 1];
    }

    has(x: T): boolean {
        if (this.size === 0) {
            return false;
        }
        const a = this.findBucket(x);
        const i = bisectLeft(a, x);
        return i !== a.length && a[i] === x;
    }

    add(x: T): boolean {
        if (this.size === 0) {
            this.a = [[x]];
            this.size = 1;
            return true;
        }
        const a = this.findBucket(x);
        const i = bisectLeft(a, x);
        if (i !== a.length && a[i] === x) {
            return false;
        }
        a.splice(i, 0, x);
        this.size++;
        if (a.length > this.a.length * SortedSet.REBUILD_RATIO) {
            this.build();
        }
        return true;
    }

    discard(x: T): boolean {
        if (this.size === 0) {
            return false;
        }
        const a = this.findBucket(x);
        const i = bisectLeft(a, x);
        if (i === a.length || a[i] !== x) {
            return false;
        }
        a.splice(i, 1);
        this.size--;
        if (a.length === 0) {
            this.build();
        }
        return true;
    }

    lt(x: T): T | null {
        for (const a of this.a.slice().reverse()) {
            if (a[0] < x) {
                return a[bisectLeft(a, x) - 1];
            }
        }
        return null;
    }

    le(x: T): T | null {
        for (const a of this.a.slice().reverse()) {
            if (a[0] <= x) {
                return a[bisectRight(a, x) - 1];
            }
        }
        return null;
    }

    gt(x: T): T | null {
        for (const a of this.a) {
            if (a[a.length - 1] > x) {
                return a[bisectRight(a, x)];
            }
        }
        return null;
    }

    ge(x: T): T | null {
        for (const a of this.a) {
            if (a[a.length - 1] >= x) {
                return a[bisectLeft(a, x)];
            }
        }
        return null;
    }

    get(x: number): T {
        if (x < 0) x += this.size;
        if (x < 0) throw new Error("IndexError");
        for (const a of this.a) {
            if (x < a.length) {
                return a[x];
            }
            x -= a.length;
        }
        throw new Error("IndexError");
    }

    index(x: T): number {
        let ans = 0;
        for (const a of this.a) {
            if (a[a.length - 1] >= x) {
                return ans + bisectLeft(a, x);
            }
            ans += a.length;
        }
        return ans;
    }

    indexRight(x: T): number {
        let ans = 0;
        for (const a of this.a) {
            if (a[a.length - 1] > x) {
                return ans + bisectRight(a, x);
            }
            ans += a.length;
        }
        return ans;
    }
}

function main() {
    const [n, k] = nextNums(2)
    const sortedSet = new SortedSet()
    // 山
    const deck = nextNums(n)
    // 場に出ているカードを管理するレコード。空の配列で初期化する
    const fieldCards: Record<number, number[]> = {}
    for (let i = 1; i < n + 1; i++) {
        fieldCards[i] = []
    }
    // 答え
    const ans: number[] = new Array(n + 1).fill(-1)
    for (let turn = 0; turn < n; turn++) {
        // カードを引く
        const card = deck[turn]
        // 今のカードより大きい最少のカードを探す
        const minCard = sortedSet.ge(card)
        // 見つかった場合
        if (minCard) {
            // 場から消す
            sortedSet.discard(minCard)
            // 場のカード情報を更新。元々積まれていたカード情報をもらうイメージ
            fieldCards[card] = fieldCards[minCard]
        }
        // 今のカードを場のカード情報に追加
        fieldCards[card].push(card)
        // 今のカードを場に出す
        sortedSet.add(card)
        // 場のカードがK以上になった場合、食べる
        if (fieldCards[card].length >= k) {
            // 中のカードを繰り返して、答えを更新
            for (const c of fieldCards[card]) {
                ans[c] = turn + 1
            }
            // 食べたため場からカードを消す
            sortedSet.discard(card)
        }
    }
    log(ans.slice(1).join(' '))
}

let inputs = "";
let inputArray: string[];
let currentIndex = 0;

function next() {
    return inputArray[currentIndex++];
}
function nextNum() {
    return +next();
}
function nextBigInt() {
    return BigInt(next());
}
function nexts(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = next();
    return arr;
}
function nextNums(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = nextNum();
    return arr;
}
function nextBigInts(length: number) {
    const arr = [];
    for (let i = 0; i < length; ++i) arr[i] = nextBigInt();
    return arr;
}


// デバッグ環境がWindowsであれば条件分岐する
if (process.env.OS == "Windows_NT") {
    const stream = createInterface({
        input: process.stdin,
        output: process.stdout,
    });
    stream.on("line", (line) => {
        inputs += line;
        inputs += "\n";
    });
    stream.on("close", () => {
        inputArray = inputs.split(/\s/);
        main();
    });
} else {
    inputs = fs.readFileSync("/dev/stdin", "utf8");
    inputArray = inputs.split(/\s/);
    main();
}


この記事が気に入ったらサポートをしてみませんか？