AtCoder精選良問「D - Knight」Diff 1009

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

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

問題リンク

D - Knight

問題概要

グリッド上の原点にナイトの駒がある

• ナイトはマス(i, j)にあるとき、(i+1, j+2)または(i+2, j+1)のどちらかのマスに動かすことができる

• ナイトの駒をマス(X, Y)まで移動させる方法は何通りあるか求める

• 答えは10^9 + 7で割った余りを出力する

• 制約：1≤X≤10^6、1≤Y≤10^6

考え方

初見ではDPっぽいと思う。
dp[i][j]  マスi,jの行き方の通り数、遷移式はdp[i][j] = dp[i-1][j-2] + dp[i-2][j-1]とすれば正解は出そう。ただ、XとYが1,000,000迄と大きいので、二次元dpで行おうとすると、計算量が大きくなりすぎてしまう。

ではどうするかと言うと、組み合わせの概念を使って解く。
縦に２マス移動した回数をA回、縦に１マス移動した回数をB回とする。
Yが縦方向のインデックスだとすると、以下が成り立つ。

① Y = 2 x A + B
次に、Xは横方向のインデックス。
縦に１マス移動した時は横に２マス進んでいるので、以下が成り立つ。
② X = A + 2 x B

XとYは入力で与えられるので、Aを全探索して、①と②が成り立つBを求める。

合計の移動回数はA + B回。
そのうち２マス移動するようなAの選び方が何通りあるか、ということを求めると答えになる。これは二項係数、つまりnCrで計算すると解ける。

なお、今回は数が大きいので、余りを取りながら計算する必要がある。余りを使った割り算は普通にやるとうまくできない。
競プロ鉄則本などで紹介されている以下の性質を使う。

Mを素数とし、bをMで割り切れない整数であるとする。このとき、Mで割った余りを求める問題では、「÷b」を「x b^M-2」に書き換えても計算結果は変わらない。

競技プログラミングの鉄則　P174

この性質を利用して、aが分子、bを分母だとすると、a * (b ^ M-2)をMで余りを取ると答えになる。b ^ M-2もこれまた計算量が大きいので、繰り返し二乗法のpower関数などを用意しておく。

実装

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

// aのb乗を高速計算する関数
export const power = (a: bigint, b: bigint, mod: bigint): bigint => {
    let ret = 1n;
    if (b === 0n) return 1n % mod
    let p = a;
    const maxBit = b.toString(2).length
    for (let i = 0n; i < maxBit; i++) {
        const wari = 2n ** i;
        if (wari > b) break
        if ((b / wari) % 2n === 1n) {
            ret = (ret * p) % mod;
        }
        p = (p * p) % mod;
    }
    return ret;
}

//a÷b を m で割った余りを返す関数
export const divisionMod = (a: bigint, b: bigint, mod: bigint): bigint => {
    const powResult = power(b, mod - 2n, mod);
    const result = (a * powResult) % mod;
    return result < 0n ? result + mod : result;
}

function main() {
    // 横、縦
    const [x, y] = nextNums(2)

    let tate2 = 0
    let tate1 = 0
    // 縦に２マス移動するのをa回行った時で全探索
    for (let a = 0; a < Math.ceil(y / 2) + 1; a++) {
        // 縦１マスで移動した回数
        const b = y - 2 * a
        // 成り立てばそれぞれの移動回数が決まる
        if (b * 2 + a == x) {
            tate2 = a
            tate1 = b
        }
    }
    const tot = BigInt(tate2 + tate1)
    // どっちも0は移動方法がない
    if (tot == 0n) {
        log('0')
        exit()
    }
    const mod = BigInt(10 ** 9 + 7)
    const r = Math.max(tate2, tate1)
    // 分子
    let numerator = 1n
    for (let i = 0n; i < r; i++) {
        numerator *= (tot - i)
        numerator %= mod
    }
    // 分母
    let denominator = 1n
    for (let i = 1n; i < r + 1; i++) {
        denominator *= i
        denominator %= mod
    }
    const ans = divisionMod(numerator, denominator, mod)
    log(String(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();
}