AtCoder精選良問「D - Index Trio」Diff 983

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

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

問題リンク

D - Index Trio

問題概要

長さNの整数列Aが与えられ、以下の条件を全て満たす整数の組(i, j, k)の総数を求める問題。

条件：

1. 1≤i, j, k≤N

2. Aj / Ai = Ak

考え方

インデックスの組み合わせを総列挙する問題
素朴に行うと３重ループで全組み合わせを試せばいいが、Nの最大の制約が200,000までと大きいので、明らかにTLEする。
そのため、何らかの高速化の手順を踏む必要がある。

まずは条件２を変形させてみる。

Aj = Ai * Ak

AjはAiとAkの積であることが分かる。

とりあえず、同じ数字をみる必要はなさそうなので、重複している数字はカットする。

たとえば以下のような配列が与えられたとする。

1 2 2 3

重複をカットすると、

1 2 3

あとそれぞれの数字の出現個数も保持しておく。

count = {1:1,2:2,3:1}

こんな感じだろうか。

これで条件を満たす組み合わせを数え上げていく。
i,j,kは重複してカウントしていいのがこの問題のポイント。つまり出現数字の個数の積だけ通り数がある。
例えば、入力が[1, 1]みたいな配列だとしても、(1,1),(1,2),(2,1),(2,2)の４つが答えになる。
上記の[1, 2, 2, 3]の例だと以下のような組み合わせと通り数になる。

1 * 1 = 1 => 1 x 1 x 1 => 1通り
1 * 2 = 2 => 1 x 2 x 2 => 4通り
1 * 3 = 3 => 1 x 1 x 1 => 1通り
2 * 1 = 2 => 2 x 1 x 2 => 4通り
3 * 1 = 2 => 1 x 1 x 2 => 2通り

次に計算量について考える。
左辺のAiとAkを全探索したらどうだろうか。
例えば最悪のケースだと1～200,000までの全ての数字がある配列。
何となく間に合わないような気もするが、実は間に合う。

ここでもう１個の制約に注目すると、配列の値の最大値も200,000。
AiとAkで全探索しても、枝刈りすれば計算量はそこまで増えない。

あらかじめソートしておいて、Ai * Akが配列の最大値を超えたらbreakする方針を取る。
試しに計算量を見積もってみる。

function main() {
    const arr: number[] = []
    // 1から200000までの配列を作る
    for (let i = 1; i < 200001; i++) {
        arr.push(i)
    }
    // 計算回数
    let cnt = 0
    // aj のループ
    for (const a of arr) {
        // akのループ
        for (const b of arr) {
            const c = a * b
            if (c > 200000) break
            cnt++
        }
    }
    log(cnt)
}

2472113

約250万回のループ。これは高速に動作する範囲内だ。

実装

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

function main() {
    const n = nextNum()
    const arr = nextNums(n)
    // 数字の出現個数のカウンタ
    const counter: Record<number, number> = {}
    // 0で初期化
    for (let i = 1; i < 200001; i++) {
        counter[i] = 0
    }
    // ユニークな数字の配列
    const uniqueNums: number[] = []
    // 数字ごとの出現個数を数える
    for (const num of arr) {
        // 出現回数が０だったらユニーク
        if (counter[num] == 0) uniqueNums.push(num)
        // カウントアップする
        counter[num]++
    }
    // 配列を昇順ソート
    uniqueNums.sort((a, b) => (a - b))
    // 配列の最大値
    const maxNum = uniqueNums[uniqueNums.length - 1]

    let ans = 0
    // aj と akを全探索
    for (const aj of uniqueNums) {
        for (const ak of uniqueNums) {
            const ai = aj * ak
            // aiが最大値を超えたらそこから先は調べる必要がない
            if (ai > maxNum) break
            // それぞれの出現個数をかける。aiが存在しなかったら0になる
            const cnt = counter[ak] * counter[aj] * counter[ai]
            ans += cnt
        }
    }
    log(ans)
}

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();
}