はまやんはまやんはまやん

hamayanhamayan's blog

Matrix Land [HackerRank Week of Code 35 D]

https://www.hackerrank.com/contests/w35/challenges/matrix-land

縦H、横Wの行列Aがある。
この行列上で以下のゲームを行う。

  • 最初は1行目の任意のセルからスタート
  • 1回の操作で左右か下に移動する
  • セルを訪れたら書いてある数をゲットし、ゲットしたら数は消える
  • H列目の任意のセルで終了する

ゲットできる点の総和の最大値は?

前提知識

解法

動的計画法で解く。

dp[y][x] := y行x列のマスにいる時の点の最大値
愚直に更新していく操作を考えると、(x,y)への遷移で考えるべきは、(0,y-1)~(W-1,y-1)からの遷移である。
 
例えば、x≦zとしてdp[y][x]からdp[y+1][z]への遷移を考える。
まず、(x,y)のマスまででdp[y][x]点取れる。
次に、そのまま下に移動して、(x,y+1)のマスへ移る。
そこから、横に移動して(x,y+1)~(z,y+1)のマスの点数を取り、dp[y+1][z]の点とする。
つまり、

rep(z, 0, W) rep(x, 0, z + 1) dp[y+1][z] = max(dp[y+1][z], dp[y][x] + [x,z]の点の総和)

をすれば遷移ができる。
 
実はこれでは最適ではない場合がある。
(x,y+1)から直接(z,y+1)に行くのではなく、(x,y+1)よりも左にちょっと行ってから、(z,y+1)に向かったほうがオトクな場合がある。
同様に(x,y+1)から(z,y+1)へ向かった後に、そこで止まらずに(z,y+1)よりも右にちょっと行ってから、(z,y+1)に戻った方がオトクな場合がある。
この二つのお得な場合を予め計算しておこう。
 
C[x] := (x,y+1)から右に向かう前に左に寄ることで手に入る点数の最大値
D[x] := (x,y+1)へ到着した後に、右に寄って戻ってくることで手に入る点数の最大値
それに加えて[x,z]の点の総和を円滑に計算できるように累積和を作っておく
B[x] := [0,x]の点数の総和
 
これを使うと更新プログラムは以下のようになる

rep(z, 0, W) rep(x, 0, z + 1) dp[y+1][z] = max(dp[y+1][z], dp[y][x] + B[z] - B[x - 1] + C[x] + D[z])

 
最大値をとっているが、よくよくみると、xが関連してくる所とそうでない所がある。
xが関連してないものを外に出すとこうなる。

rep(z, 0, W) {
    int ma = -INF;
    rep(x, 0, z + 1) ma = max(ma, dp[y][x] - B[x - 1] + C[x])
    dp[y+1][z] = max(dp[y+1][z], ma + B[z] + D[z]);
}

 
すると最大を探している部分はセグメントツリーで高速化できることが分かる。
そのため、区間最大のセグメントツリーにこの部分を置いて高速に取得できるようにすれば計算量的に間に合う。

int H, W, A[4010101], _B[4010101], C[4010101], D[4010101];
int *B;
SegTree<int, 1 << 22> st;
//---------------------------------------------------------------------------------------------------
void _main() {
    make_test();
    return;

    cin >> H >> W;
    B = _B + 1;

    vector<int> dp(W, 0);
    rep(y, 0, H) {
        rep(x, 0, W) cin >> A[x];
        vector<int> pd(W, -INF);

        rep(_, 0, 2) {
            rep(x, 0, W) B[x] = A[x];
            rep(x, 1, W) B[x] += B[x - 1];
            
            int ma = 0;
            rep(x, 0, W) {
                ma += A[x];
                if (ma < 0) ma = 0;
                C[x + 1] = ma;
            }

            ma = 0;
            rrep(x, W - 1, 1) {
                ma += A[x];
                if (ma < 0) ma = 0;
                D[x - 1] = ma;
            }

            rep(x, 0, W) st.update(x, dp[x] - B[x - 1] + C[x]);

            rep(x, 0, W) {
                int ma = st.get(0, x + 1);
                pd[x] = max(pd[x], ma + B[x] + D[x]);
            }

            reverse(dp.begin(), dp.end());
            reverse(pd.begin(), pd.end());
            reverse(A, A + W);
        }

        swap(dp, pd);
    }

    int ans = -INF;
    rep(x, 0, W) ans = max(ans, dp[x]);
    cout << ans << endl;
}