見出し画像

AtCoder精選良問「D - Practical Skill Test」Diff 1220

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

問題リンク

問題概要

まず、H,W,Dが入力で与えられる。
H x W行のマスの中は1からH x Wまでの数字がランダムで振られている。
その後、クエリでLとRが与えられる。
最初駒はマスの中のLで与えられた数字がある場所に位置している。
その後、L から+ Dの移動を繰り返してRに達した時にマンハッタン距離で合計いくら動いたか、ということを各クエリごとに出力する、という問題

考え方

まず、「L から+ Dの移動を繰り返してRに達した時にマンハッタン距離で合計いくら動いたか」ということがどういうことなのか、入力例からみてみる。
H=3, W=3, D=2の場合。

1 4 3
2 5 7
8 9 6

このようなマスが与えられている。

制約からR-LはDの倍数になる且つ、LはRより小さい。
例えばL=2,R=4みたいな入力が考えられる。
2は上のマスでのゼロインデックスのポジションで表すと(1,0)。
4は(0,1)となる。そのため、マンハッタン距離は|1-0| + |0-1|で2だ。
これを順番に答えていくということになる。

次に制約に注目をしてみる。
HとWの制約は300以内のため、数字は最大で90,000。
クエリの数は最大で10**5ほどと大きい。
そのため、LとRが与えられたときに愚直な計算をしていくと、当然TLEをする。LとRの数字から瞬時に答えを出すことが求められそうな問題だ。

この手の問題では累積和のアルゴリズムが使えそう。
例えば2→4→6→8という移動距離の累積和を事前に計算しておく。
まず各遷移の移動距離は上のマスから以下のようになる。

// 2→4→6→8の移動のマンハッタン距離
[0, 2, 3, 2]

そのため累積和は以下のようになる。


// 2→4→6→8の移動のマンハッタン距離の累積和
[0, 2, 5, 7]

このような配列を事前に持っていれば、例えば4→8というクエリが来た際は、7-2で5が答えだな、ということがO(1)で計算することができる。
クエリは4→8ではなくて1→3というように来ることもあるので、各移動パターンの累積和を事前に用意しておく、ということが問題を解くカギになりそう。

実装の流れ

まずは、[1,3,5,7,9],[2,4,6,8]みたいな配列を列挙する方法を考える。
色々と方法はあると思うが、数字はDずつ増えていくので、二重ループでDずつインクリメントしていく。

function main() {
    const [h, w, d] = nextNums(3)
    // [[1,3,5,7],[2,4,6,8]]みたいな配列
    const indexArr: number[][] = []
    const maxNum = h * w
    for (let i = 1; i < d + 1; i++) {
        const arr: number[] = []
        arr.push(i)
        for (let j = i + d; j <= maxNum; j += d) {
            arr.push(j)
        }
        indexArr.push(arr)
    }
    log(indexArr)
}
// 出力
[ [ 1, 3, 5, 7, 9 ], [ 2, 4, 6, 8 ] ]

次にこの配列に従って移動した際のマンハッタン距離を算出していく。
まず数字ごとのポジションを記録したレコードを持っておくといい。

// {1:[0,0],2:[1,0]}というように数字ごとのポジションを記録
const posDict: Record<number, number[]> = {}
for (let i = 0; i < h; i++) {
    for (let j = 0; j < w; j++) {
        const num = nextNum()
        posDict[num] = [i, j]
    }
}
log(posDict)
{
  '1': [ 0, 0 ],
  '2': [ 1, 0 ],
  '3': [ 0, 2 ],
  '4': [ 0, 1 ],
  '5': [ 1, 1 ],
  '6': [ 2, 2 ],
  '7': [ 1, 2 ],
  '8': [ 2, 0 ],
  '9': [ 2, 1 ]
}

良い感じ。
あとはindexArrを繰り返してマンハッタン距離を保存しておく。

// マンハッタン距離を保存
const manDistArr: number[][] = []
for (const arr of indexArr) {
    const tmp = []
    // 起点は0
    tmp.push(0)
    for (let j = 1; j < arr.length; j++) {
        // 一つ前のポジションと今のポジションを取得
        const prevNum = arr[j - 1]
        const nowNum = arr[j]
        const prevPos = posDict[prevNum]
        const nowPos = posDict[nowNum]
        // マンハッタン距離の算出
        const dist = Math.abs(prevPos[0] - nowPos[0]) + Math.abs(prevPos[1] - nowPos[1])
        // 一つ前と足して累積和化しておく
        tmp.push(dist+tmp[j-1])
    }
    manDistArr.push(tmp)
}
log(manDistArr)
[ [ 0, 2, 4, 5, 7 ], [ 0, 2, 5, 7 ] ]

良い感じに累積和の配列にできた。
あとはこれを元にクエリに答えていくということになる。

最後に問題になりそうなのは作った配列のどこを使って計算するか、ということ。
もう一度indexArrとmanDistArrをみてみる。

// indexArr
[ [ 1, 3, 5, 7, 9 ], [ 2, 4, 6, 8 ] ]
// manDistArr
[ [ 0, 2, 4, 5, 7 ], [ 0, 2, 5, 7 ] ]

indexArrはdで割った余りの順番に並んでいることが分かる。
つまり、余り1,余り2,,,余りd-1,余り0という風に並ぶ。
今回はdが2だったので少し分かりづらいかもしれない。
例えばh=5,w=5,d=4の時は以下のようになる。

[[1,  5,  9, 13, 17, 21, 25],
[ 2, 6, 10, 14, 18, 22 ],
[ 3, 7, 11, 15, 19, 23 ],
[ 4, 8, 12, 16, 20, 24 ]]

そのため、縦のインデックスについては、LもしくはRをDで割った余りを計算すれば特定ができる。

次に横のインデックスは例えばPythonならbisectLeftなどの二分探索モジュールを使えば高速に特定ができる。
残念ながらTypeScriptにはないので、自前で関数を用意する。

/**
 * 
 * @param arr ソート済みの配列
 * @param x 検索値
 * @returns 
 * 値 x が最初に現れるインデックスを返す。
 * x が arr 内に存在しない場合は、x を挿入するためのインデックスを返す
 */
export const bisectLeft = (arr: number[], x: number, lo: number = 0, hi: number = arr.length): number => {
    while (lo < hi) {
        const mid = Math.floor((lo + hi) / 2);
        if (arr[mid] < x) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    return lo;
}

以上のことから各クエリの計算部分は以下のようになる。

const q = nextNum()
for (let i = 0; i < q; i++) {
    const [l, r] = nextNums(2)
    // 縦のインデックスを特定
    let rowIdx = l % d - 1
    // 余り0の場合は最後
    if (rowIdx == -1) rowIdx = indexArr.length - 1
    // 該当するインデックスの配列を取得
    const arr = indexArr[rowIdx]
    // 対応した累積和の配列を取得
    const distArr = manDistArr[rowIdx]
    // l,rの値から横のインデックスを取得
    const left = bisectLeft(arr, l)
    const right = bisectLeft(arr, r)
    const ans = distArr[right] - distArr[left]
    log(ans)
}

コード全文

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

/**
 * 
 * @param arr ソート済みの配列
 * @param x 検索値
 * @returns 
 * 値 x が最初に現れるインデックスを返す。
 * x が arr 内に存在しない場合は、x を挿入するためのインデックスを返す
 */
export const bisectLeft = (arr: number[], x: number, lo: number = 0, hi: number = arr.length): number => {
    while (lo < hi) {
        const mid = Math.floor((lo + hi) / 2);
        if (arr[mid] < x) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    return lo;
}

function main() {
    const [h, w, d] = nextNums(3)
    // [[1,3,5,7],[2,4,6,8]]みたいな配列
    const indexArr: number[][] = []
    const maxNum = h * w
    for (let i = 1; i < d + 1; i++) {
        const arr: number[] = []
        arr.push(i)
        for (let j = i + d; j <= maxNum; j += d) {
            arr.push(j)
        }
        indexArr.push(arr)
    }

    // {1:[0,0],2:[1,0]}というように数字ごとのポジションを記録
    const posDict: Record<number, number[]> = {}
    for (let i = 0; i < h; i++) {
        for (let j = 0; j < w; j++) {
            const num = nextNum()
            posDict[num] = [i, j]
        }
    }

    // マンハッタン距離を保存
    const manDistArr: number[][] = []
    for (const arr of indexArr) {
        const tmp = []
        // 起点は0
        tmp.push(0)
        for (let j = 1; j < arr.length; j++) {
            // 一つ前のポジションと今のポジションを取得
            const prevNum = arr[j - 1]
            const nowNum = arr[j]
            const prevPos = posDict[prevNum]
            const nowPos = posDict[nowNum]
            // マンハッタン距離の算出
            const dist = Math.abs(prevPos[0] - nowPos[0]) + Math.abs(prevPos[1] - nowPos[1])
            // 一つ前と足して累積和化しておく
            tmp.push(dist + tmp[j - 1])
        }
        manDistArr.push(tmp)
    }
    const q = nextNum()
    for (let i = 0; i < q; i++) {
        const [l, r] = nextNums(2)
        // 縦のインデックスを特定
        let rowIdx = l % d - 1
        // 余り0の場合は最後
        if (rowIdx == -1) rowIdx = indexArr.length - 1
        // 該当するインデックスの配列を取得
        const arr = indexArr[rowIdx]
        // 対応した累積和の配列を取得
        const distArr = manDistArr[rowIdx]
        // l,rの値から横のインデックスを取得
        const left = bisectLeft(arr, l)
        const right = bisectLeft(arr, r)
        const ans = distArr[right] - distArr[left]
        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();
}

この記事が気に入ったらサポートをしてみませんか?