AtCoder Beginner Contest 027 D - ロボット

競技プログラミング 動的計画法

問題

https://beta.atcoder.jp/contests/abc027/tasks/abc027_d

数直線の原点にロボットが置かれている。 はじめ、ロボットの幸福度は 0 である。 このロボットに命令列が与えられる。 命令列は次の 3 文字のみからなり、先頭から末尾まで順に実行される。

M : 正または負の好きな向きに、距離 1 だけ移動する。 + : 今の座標を x とすると、幸福度が +x だけ変化する。 - : 今の座標を x とすると、幸福度が −x だけ変化する。

命令列を実行し終えた後、 ロボットは原点に戻っていなければならない 。 命令列を実行している間、ロボットの座標および幸福度は負になり得る。 最終的な幸福度を最大化するようにロボットが移動したとき、最終的な幸福度を求めよ。

方針

部分点解法であるDPで解くこととします。計算量は O(N^2)

dp[i][j] := i 回目の命令で j の位置にいる時の最大値とします。( i 回目の命令= i 文字目の文字とします。)

dp[i][j]のときに、dp[i+1][j]への更新パターンは以下のように考えられます。

  • 命令 i がM の場合
    i+1 回目に位置 j にいるためには、i 回目に位置 j+1 または j-1 にいる必要があります。

dp[i+1][j] = dp[i][j+1] または dp[i][j-1] の大きい方

  • 命令 i が + の場合
    i 回目に位置 j にいるとき、i+1 回目は j 分だけ値が増えます。

dp[i+1][j] = dp[i][j] + j

  • 命令 i が - の場合
    i 回目に位置 j にいるとき、i+1 回目は j 分だけ値が減ります。

dp[i+1][j] = dp[i][j] - j

以上の場合をまとめればいいのですが、地点 j に関しては負の値も取りうるので、
基準点を n 分だけずらしておきます。
また、値の開始位置は n になるため、dp[0][n]の値は 0、他は -INF に初期化しておきます。
最終的に一番最初にいた位置(地点 n )にいる必要があるので、求めるべき解はdp[n][n] です。

実装

import java.util.Arrays;
import java.util.Scanner;

public class Main {

    static int INF = 1 << 30;

    public static void main(String[] args) {

        Scanner in = new Scanner(System.in);
        char[] s = in.next().toCharArray();
        int n = s.length;

        int offset = n;
        int[][] dp = new int[n+1][n+1+offset];

        for (int i = 0; i <= n; i++) {
            Arrays.fill(dp[i], -INF);
        }

        dp[0][offset] = 0;

        for (int i = 0; i < n; i++) {
            for (int j = 1; j < n+offset; j++) {
                if (s[i] == 'M') {
                    dp[i+1][j] = Math.max(dp[i][j+1], dp[i][j-1]);
                } else if (s[i] == '+') {
                    dp[i+1][j] = dp[i][j] + (j-offset);
                } else if (s[i] == '-') {
                    dp[i+1][j] = dp[i][j] - (j-offset);
                }
            }
        }

        System.out.println(dp[n][offset]);
    }
}