PKU 1036 Gangsters

Gangsters

動的計画法。O(NKT)は間に合わなかった。

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main()
{
    int N, K, T;  cin >> N >> K >> T;
    vector<int> t(N), P(N), S(N);
    for ( int i=0; i<N; i++ )  cin >> t[i];
    for ( int i=0; i<N; i++ )  cin >> P[i];
    for ( int i=0; i<N; i++ )  cin >> S[i];

    vector<vector<int> > gang(T+1);
    for ( int i=0; i<N; i++ )
    {
        if ( gang[t[i]].empty() )
            gang[t[i]] = vector<int>(K+1);
        gang[t[i]][S[i]] += P[i];
    }

    vector<int> D(K+1,-99999);
    D[0] = 0;
    for ( int i=1; i<=T; i++ )
    {
        vector<int> tmp = D;
        for ( int j=0; j<=K; j++ )
        {
            D[j] = tmp[j];
            if ( j-1 >= 0 )  D[j] = max( D[j], tmp[j-1] );
            if ( j+1 <= K )  D[j] = max( D[j], tmp[j+1] );

            if ( ! gang[i].empty() )
                D[j] += gang[i][j];
        }
    }

    cout << *max_element(D.begin(),D.end()) << endl;    
}