AtCoder精選良問「D - Preparing Boxes」Diff 926

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

AtCoder 精選良問集 Diff 茶後半～水前半｜Qlitre｜note

問題リンク

D - Preparing Boxes

問題概要

空の箱がN個横一列に並んでおり、左からi番目の箱には整数iが書かれている。すぬけさんは各箱にボールを1個入れるか、何も入れないかを選べる。以下の条件を満たすボールの入れ方を「いいボールの入れ方」と定める。

• 1以上N以下の任意の整数iについて、iの倍数が書かれた箱に入っているボールの個数の和を2で割った余りがaiである。

「いいボールの入れ方」が存在するかどうかを判定し、存在する場合は1つ求めて以下を出力する。

• ボールを入れた箱の数

• ボールを入れた箱の番号

存在しない場合は-1を出力する。

考え方

解いた後で面白い、と思った。
まず問題文の「1以上N以下の任意の整数iについて、iの倍数が書かれた箱」という部分を整理してみる。
例えば１～６までの箱がある場合。
１の倍数：１，２，３，４，５，６
２の倍数：２，４，６
３の倍数：３，６
４の倍数：４
５の倍数：５
６の倍数：６

というような形になる。

これを見てみると、前から番号を追っていくと難しいことが分かる。
例えば６の２と３の倍数なので、２の時に正しいと思って６にボールを入れたけど、３で見るとボールを入れるのは間違いだったというようなことが起きる。

こういう時は発想を転換させて「後ろから考える」という手法を使う。
大きい数字から順番に追っていけば、操作の矛盾が起きない。

aiのリストは仮にこういう風にしておこう。

0 1 0 0 0 1

まずは、６にボールを入れる必要があるので、ボールを入れる。
５と４は0なので入れない。
３を見た時に3と6の合計を２で割った余りを0にする必要があるので３にもボールを入れる。
２を見た時は同様に余りを１にする必要がある。６に１個だけボールが入っているので、何もしなくていい。

このような感じで処理を追っていく。

次に計算量だが、箱のサイズの最大は200,000。
愚直に倍数をたどっていくと、TLEしそうな気配があるが意外とそうでもない。
全部繰り返さないといけないのは１の時だけで、それ以外の数字の時は数が大きいとそれだけ繰り返しが減る。

２のときは２，４，６，８…
３のときは３，６，９，１２…

この辺は事前に検証してみてもいいだろう。

function main() {
    const n = 200000
    let cnt = 0
    for (let i = 1; i <= n; i++) {
        // nをiを割った商がだいたいの計算量
        cnt += Math.floor(n / i)
    }
    log(cnt)
}

2472113

このように最大値で250万回程度のループで終えられる。
あとは実装をすればいい。

実装

上のルールで行っていけば、必ず「いいボールの入れ方」となるので、答えが-1になることはない。

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

function main() {
    const n = nextNum()
    // 1インデックスにしておく。
    const arr = [0, ...nextNums(n)]
    // ボールを入れた回数
    let countOperation: number = 0
    // ボールを入れた箱
    const boxWithBalls: number[] = []
    // 箱の状態を管理　これも１インデックスに対応
    const boxes: number[] = new Array(n + 1).fill(0)
    // 後ろから最初まで
    for (let i = n; i > 0; i--) {
        let totalBall: number = 0
        let p = 2
        // iの倍数を順番に見て入れたボールの個数を計算
        while (i * p <= n) {
            totalBall += boxes[i * p]
            p += 1
        }
        // 偶奇が違う場合は箱に入れる
        if (totalBall % 2 != arr[i]) {
            boxes[i] = 1
            countOperation++
            boxWithBalls.push(i)
        }
    }
    log(countOperation)
    boxWithBalls.reverse()
    log(boxWithBalls.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();
}