【牛客编程巅峰赛S1第4场】寻找 520

题目

又到了一年 5.20！牛牛准备了一个字符串 S，里面全都是由 '5'、'2' 和 '0' 组成，但有的数字被油污染了，看不清，以 '?' 表示。
如果把这些被油污染的部分换成 '5'、'2' 和 '0'，在所有可能的字符串中，位置不同的为 "520" 的子序列共有多少个？
将答案模 998244353 返回。

解题思路

动态规划：

使用 $f[i][0]$ 表示前 $i$ 个字符中 '5' 的个数。
使用 $f[i][1]$ 表示前 $i$ 个字符中，位置不同的为 "52" 的子序列的总个数。
使用 $f[i][2]$ 表示前 $i$ 个字符中，位置不同的为 "520" 的子序列的总个数。

状态转移方程：

当 $S[i] = 5$ 时， $f[i][0] = f[i-1][0] + 1$ ，否则 $f[i][0] = f[i-1][0]$ 。
当 $S[i] = 2$ 时，当前的字符 '2' 可以和前 $i-1$ 个字符中的 $f[i-1][0]$ 个 '5' 组成 $f[i-1][0]$ 个 "52" 子序列。
所以， $f[i][1] = f[i-1][1] + f[i-1][0]$ 。
当 $S[i] = 0$ 时，当前的字符 '0' 可以和前 $i-1$ 个字符中的 $f[i-1][1]$ 个 "52" 子序列组成 $f[i-1][1]$ 个 "520" 子序列。
所以， $f[i][2] = f[i-1][2] + f[i-1][1]$ 。
当 $S[i] = ?$ 时，
如果当前的 $?$ 是 '0'，那么 $f[i][2] = f[i-1][2] + f[i-1][1]$ ；
如果当前的 $?$ 是 '2'，那么 $f[i][1] = f[i-1][1] + f[i-1][0]$ ；
如果当前的 $?$ 是 '5'，那么 $f[i][0] = f[i-1][0] + 1$ 。

此时，时间复杂度 $O(n)$ ，空间复杂度 $O(n)$ 。 $n$ 为字符串 $S$ 的长度。
根据上面的状态转移方程，可以看出 $f[i][0]$ 、 $f[i][1]$ 和 $f[i][2]$ 只与它们的前一个状态相关。
所以，可以对上面的 dp 进行空间上的优化。代码如下。
此时空间复杂度为 $O(1)$ 。

C++代码

class Solution {
public:
    /**
     * 寻找所有的 520 子序列
     * @param S string字符串 牛牛准备的数字字符串
     * @return int整型
     */
    int findOccurrences(string S) {
        // write code here
        const int mod = 998244353;
        int f[3] = {0};
        for(int i=0; i<S.size(); ++i){
            if(S[i]=='5')
                f[0] += 1;
            else if(S[i]=='2')
                f[1] += f[0];
            else if(S[i]=='0')
                f[2] += f[1];
            else{
                f[2] += f[1];
                f[1] += f[0];
                f[0] += 1;
            }
            f[0] %= mod;
            f[1] %= mod;
            f[2] %= mod;
        }
        return f[2];
    }
};