AtCoder精選良問「D - Index Trio」Diff 983
この記事はAtCoderの精選良問集マガジンの中の記事です。
私が良問だと思った問題を解説付きでまとめています。
問題リンク
問題概要
長さNの整数列Aが与えられ、以下の条件を全て満たす整数の組(i, j, k)の総数を求める問題。
条件：
1≤i, j, k≤N
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();
}