ロボットがグリッド内を移動する際の総コストを求めるC++プログラムの解説

問題の概要

縦 h × 横 w のグリッドが与えられ、各マスには正の整数が書かれています。このグリッド上には経路探索ロボットが配置されており、あるマス (p, q)(p は行番号、q は列番号)から別のマス (i, j) へ移動できます。1 回の移動にかかるコストは |p − i| + |q − j| として定義されます。

さらに、次のような性質を持つ q 回の移動(トリップ)が与えられます。

• 各トリップは 2 つの値 (x, y) を持ち、すべてのトリップで共通の値 d が使われます。
• ロボットは値 x が書かれたマスに出発し、値 x + d が書かれた別のマスへ移動します。
• その後も値 x + d + d、x + d + d + d … という具合に移動を続け、値が y 以上のマスに到達した時点で終了します。
• y − x は必ず d の倍数になることが保証されています。

それぞれのトリップについて合計コストを求めてください。ただし、ロボットが一度も移動できない場合は、そのトリップのコストは 0 となります。

入出力例

例として、h = 3、w = 3、d = 3、q = 1、grid = {{2, 6, 8}, {7, 3, 4}, {5, 1, 9}}、trips = {{3, 9}} が与えられた場合を考えます。このとき出力は 4 になります。

• 3 はマス (2, 2) に存在します
• 6 はマス (1, 2) に存在します
• 9 はマス (3, 3) に存在します

(2, 2) → (1, 2) の移動コストは |2 − 1| + |2 − 2| = 1、(1, 2) → (3, 3) の移動コストは |1 − 3| + |2 − 3| = 3 となるため、合計コストは 1 + 3 = 4 です。

解法のアプローチ

この問題は、次の手順で効率よく解くことができます。

1. まず、map を使って「各値がどの座標にあるか」を事前に記録します。
2. 次に、d ごとに「値がちょうど d だけ異なるマス間の移動コスト」を計算し、累積和(プレフィックスサム)として配列 dp に格納します。
3. 最後に、各トリップに対して累積和を参照すれば、区間内のコスト合計を O(1) で算出できます。

具体的な流れは、以下の擬似コードの通りです。

マップ loc を定義する
i := 0 から開始し、i < h の間 i を 1 ずつ増やしながら繰り返す:
    j := 0 から開始し、j < w の間 j を 1 ずつ増やしながら繰り返す:
        loc[grid[i, j]] := 新しいペア (i, j)
配列 dp[d + 1] を定義する
i := 1 から開始し、i <= d の間 i を 1 ずつ増やしながら繰り返す:
    j := i とする
    j < w * h の間、以下を繰り返す:
        n := j + d とする
        もし j + d > w * h ならば:
            ループを抜ける
        dx := |loc[n] の 1 番目の要素 − loc[j] の 1 番目の要素|
        dy := |loc[n] の 2 番目の要素 − loc[j] の 2 番目の要素|
        j := j + d とする
        dp[i] の末尾に dx + dy を追加する
    j := 1 から開始し、j < dp[i] のサイズの間 j を 1 ずつ増やしながら繰り返す:
        dp[i, j] := dp[i, j] + dp[i, j − 1](累積和を作成)
i := 0 から開始し、i < q の間 i を 1 ずつ増やしながら繰り返す:
    tot := 0 とする
    le := trips[i] の 1 番目の値
    ri := trips[i] の 2 番目の値
    もし ri mod d が 0 ならば:
        f := d
    そうでなければ:
        f := ri mod d
    pxl := (le − f) / d
    pxr := (ri − f) / d
    もし le が f と等しければ:
        もし ri が f と等しければ:
            tot := 0
        そうでなければ:
            tot := tot + (dp[f, pxr − 1] − 0)
    そうでなければ:
        もし ri が f と等しければ:
            tot := 0
        そうでなければ:
            tot := tot + dp[f, pxr − 1] − dp[f, pxl − 1]
    tot を出力する

この手法では、前計算に O(h × w)、各クエリへの回答に O(1) しかかからないため、クエリ数が非常に多いケースでも高速に処理できるのが大きな強みです。

C++ 実装例

理解を深めるために、以下の実装例を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
const int INF = 1e9;
void solve(int h, int w, int d, int q, vector<vector<int>> grid,
vector<pair<int, int>> trips) {
    map<int, pair<int, int>> loc;
    for (int i = 0; i < h; i++) {
        for (int j = 0; j < w; j++)
            loc[grid[i][j]] = make_pair(i, j);
    }
    vector<int> dp[d + 1];
    for (int i = 1; i <= d; i++) {
        int j = i;
        while (j < w * h) {
            int n = j + d;
            if (j + d > w * h)
                break;
            int dx = abs(loc[n].first - loc[j].first);
            int dy = abs(loc[n].second - loc[j].second);
            j += d;
            dp[i].push_back(dx + dy);
        }
        for (j = 1; j < dp[i].size(); j++)
            dp[i][j] += dp[i][j - 1];
    }
    for (int i = 0; i < q; i++) {
        int tot = 0;
        int le, ri;
        le = trips[i].first;
        ri = trips[i].second;
        int f;
        if (ri % d == 0)
            f = d;
        else
            f = ri % d;
        int pxl, pxr;
        pxl = (le - f) / d;
        pxr = (ri - f) / d;
        if (le == f){
            if (ri == f)
                tot = 0;
            else
                tot += (dp[f][pxr - 1] - 0);
        } else {
            if (ri == f)
                tot = 0;
            else
                tot += dp[f][pxr - 1] - dp[f][pxl - 1];
        }
        cout<< tot << endl;
    }
}
int main() {
    int h = 3, w = 3, d = 3, q = 1;
    vector<vector<int>> grid = {{2, 6, 8}, {7, 3, 4}, {5, 1, 9}};
    vector<pair<int, int>> trips = {{3, 9}};
    solve(h, w, d, q, grid, trips);
    return 0;
}

入力

3, 3, 3, 1, {{2, 6, 8}, {7, 3, 4}, {5, 1, 9}}, {{3, 9}}

出力

4