yukicoder No.258 回転寿司(2)

頭の悪い方法で解いてしまった。writerの解説(http://yukicoder.me/problems/710/editorial)のほうが計算量的に優れています。

問題

No.258 回転寿司(2) - yukicoder

解法

メモ化再帰で解きました。cur番目の寿司を食べるか食べないか決める時,curを食べずにcur+1を食べるのが最適な場合はcurの次に食べる寿司ne[cur]をcur+1とし,curを食べてcur+2以降を食べるのが最適な場合はne[cur] = i(iはcur+2以降の何らかの数)としてdp[cur]とne[cur]を埋めていきます。

この情報を用いて経路復元をやっていきます。

int dp[1011];
int ne[1011];

int n;
int V[1011];

int dfs(int cur) {
    if (cur >= n) return 0;
    int& ret = dp[cur];
    if (ret >= 0) return ret;
    ret = dfs(cur+1);
    for (int i = cur+2; i <= n+1; i++) {
        if (ret < dfs(i)+V[cur]) {
            ne[cur] = i;
            ret = dfs(i)+V[cur];
        }
    }
    return ret;
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    cin >> n;
    for (int i = 0; i < n; i++) cin >> V[i];
    memset(dp, -1, sizeof(dp));
    for (int i = 0; i < n; i++) ne[i] = i+1;
    int ans = dfs(0);
    cout << ans << endl;
    vector<int> select;
    int cur = 0;
    while (cur < n) {
        if (ne[cur] == cur+1) cur = cur+1;
        else {
            select.push_back(cur+1);
            cur = ne[cur];
        }
    }
    int m = select.size();
    for (int i = 0; i < m; i++) {
        cout << select[i] ;
        if (i < m-1) cout << " ";
    }
    cout << endl;
    return 0;
}