AtCoder精選良問「D - Handstand 2」Diff 1045

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

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

問題リンク

D - Handstand 2

問題概要

正の整数Nが与えられ、N以下の組(A,B)で、以下の条件を満たすものの個数を求める問題。

• A,Bは先頭に0のつかない10進数表記で表される。

• Aの末尾の桁とBの先頭の桁が等しい。

• Aの先頭の桁とBの末尾の桁が等しい。

考え方

シンプルであるものの、難しい。
まず愚直に思いつく方法としては、A,Bの組み合わせを全探索すること。
だけど、Nの制約が200,000まであるので、明らかにTLEする。

初見ではAを全探索して、Bの数を数え上げるという方法をとった。だが、これがかなりややこしい。

• Aの最初と最後の数字を取り出す

• Nを超えない範囲でBがいくつあるか数えあげる。

Bの数え上げがかなり複雑になってしまう。
例えば入力例で、N=2020の場合。Aが32の時はBはこんな感じになる。

1. 23

2. 2〇3

3. 2〇〇3

1と2の時は何となく数え上げられそうだが、桁がNと同じ時は難しい。
余談だが、初見ではこの方法で数え上げてACしたものの、かなりきつかった。

Submission #41194617 - AtCoder Beginner Contest 152

たぶん、もっと簡単な方法がある。
ACをした後に公式で紹介されているはまやんさんの解説を見て、目からうろこが落ちる思いをした。

Handstand 2 [AtCoder Beginner Contest 152 D] - はまやんはまやんはまやん

１からNまでの数字を全探索して、cnt[最初の桁][最後の桁]という表を用意して、カウントアップしていけば簡単に数が出る、とのこと。
確かにこうすれば、N以下で、先頭と末尾がある数字で固定されている数字の個数が簡単に分かる。

例えば、N=25の場合。表の中で最初の数字が1で最後の数字が2の数字が1つ（12）あり、最初の数字が2で最後の数字が1の数字が1つ（21）あることが分かる。
これらの数字の組み合わせは1つしかない。（12と21）。

同じ数字の場合はどうだろう。
例えば最初の数字が1で最後の数字が1の数字は2つ（1、11）ある。
これらの数字の組み合わせは、「最初の数字が1で最後の数字が1の数字の数」（2つ）と「最後の数字が1で最初の数字が1の数字の数」（2つ）の積、つまり2 * 2 = 4つある。
そういわれてみるとそうだ。A,Bの世界で考えた時に(1,1)、(1,11)、(11,1)、(11,11)の４つの組み合わせがある。

一般化すると、以下の方法で答えが出せる

• 1からNまで繰り返してcnt[最初の桁][最後の桁]をカウントアップする

• 1から9までの二重の繰り返しでcnt[i][j] * cnt[j][i]を足し上げる。

必死になって解いた問題で、もっと簡単な解法が見つかることがある。
これが競技プログラミングの醍醐味の一つである。

実装

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

function main() {
    const n = nextNum()
    // cnt配列を初期化
    const cnt: number[][] = []
    for (let i = 0; i < 10; i++) {
        const arr: number[] = new Array(10).fill(0)
        cnt.push(arr)
    }
    // 数え上げる
    for (let i = 1; i < n + 1; i++) {
        const firstDigit = Number(String(i)[0])
        const lastDigit = i % 10
        // 末尾がゼロは飛ばす
        if (lastDigit == 0) continue
        // 個数をインクリメント
        cnt[firstDigit][lastDigit]++
    }
    let ans = 0
    for (let i = 1; i < 10; i++) {
        for (let j = 1; j < 10; j++) {
            // 先頭がiで末尾がjの数と、先頭がjで末尾がiの数の積を足す
            ans += cnt[i][j] * cnt[j][i]
        }
    }
    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();
}