mayoko’s diary

プロコンとかいろいろ。

SRM 495 div1 easy:ColorfulCards

解法

各カードについて, あり得る数字がどれだけあるかを列挙することを考えます。
m = colors.size() としましょう。で,

ok1[i][j] = (i 番目のカードが j という値をとる可能性があるかを 0 〜 i 番目のカードのみで考える)
ok2[i][j] = (i 番目のカードが j という値をとる可能性があるかを i 〜 m-1 番目のカードのみで考える)

という dp のようなものを考えます。
例えば問題の一番最初のサンプル N = 5, colors = "RRR" を考えてみると,

ok1[0][2] = true, ok1[0][3] = true, ok1[0][5] = true となります。0 番目のカードより右側のカードの事情は考えないので, こういうことになるわけです。
一方で, ok2[0][2] = true, ok2[0][3] = false, ok2[0][5] = false となります。今度は右側のカードの事情も考えないと行けないので, 最初から大きい素数を使ってしまうとダメになってしまった, ということですね。

こんな感じの true, false を決定できれば, ok1[i][j] = true かつ ok2[i][j] = true のときに i 番目のカードは j を取りうる, と結論付けることが出来ます。

ということで, あとは ok1, ok2 を求める方法を考えれば良いですが, 例えば ok1[i][j] は, まず i 番目のカードが j が素数/合成数であるという条件と合っているなら第一条件はクリアで, あとは i-1 番目の数を 1, 2, ..., j-1 のどれかにしても問題なければ i 番目の数は j で大丈夫, となります。これはメモ化再帰的に書けますね。ok2 も同様です。

bool prime[1010];
int n, m;
int ok1[55][1010], ok2[55][1010];
string colors;

int dfs1(int index, int num) {
    int& ret = ok1[index][num];
    if (ret >= 0) return ret;
    if (colors[index] == 'R') {
        if (prime[num]) {
            if (index == 0) return ret = 1;
            for (int i = 1; i < num; i++) {
                if (dfs1(index-1, i)) return ret = 1;
            }
            return ret = 0;
        } else return ret = 0;
    } else {
        if (prime[num]) return ret = 0;
        else {
            if (index == 0) return ret = 1;
            for (int i = 1; i < num; i++) {
                if (dfs1(index-1, i)) return ret = 1;
            }
            return ret = 0;
        }
    }
}

int dfs2(int index, int num) {
    int& ret = ok2[index][num];
    if (ret >= 0) return ret;
    if (colors[index] == 'R') {
        if (prime[num]) {
            if (index == m-1) return ret = 1;
            for (int i = num+1; i <= n; i++) {
                if (dfs2(index+1, i)) return ret = 1;
            }
            return ret = 0;
        } else return ret = 0;
    } else {
        if (prime[num]) return ret = 0;
        else {
            if (index == m-1) return ret = 1;
            for (int i = num+1; i <= n; i++) {
                if (dfs2(index+1, i)) return ret = 1;
            }
            return ret = 0;
        }
    }
}

class ColorfulCards {
public:
    vector <int> theCards(int N, string COLORS) {
        for (int i = 2; i < 1010; i++) {
            prime[i] = true;
        }
        for (int i = 2; i < 1010; i++) {
            if (prime[i]) {
                for (int j = 2; i*j < 1010; j++) prime[i*j] = false;
            }
        }
        n = N;
        colors = COLORS;
        m = colors.size();
        vector<vi> nums(m);
        memset(ok1, -1, sizeof(ok1));
        memset(ok2, -1, sizeof(ok2));
        for (int i = 0; i < m; i++) {
            for (int j = 1; j <= N; j++) {
                if (!dfs1(i, j)) continue;
                if (!dfs2(i, j)) continue;
                nums[i].push_back(j);
            }
        }
        vector<int> ret(m);
        for (int i = 0; i < m; i++) {
            assert(nums[i].size() != 0);
            if (nums[i].size() > 1) ret[i] = -1;
            else ret[i] = nums[i][0];
        }
        return ret;
    }
};