牛客小白月赛43

E - 满意的集合

https://ac.nowcoder.com/acm/contest/11220/E

时间限制：C/C++ 1秒，其他语言2秒
空间限制：C/C++ 262144K，其他语言524288K
64bit IO Format: %lld

题目描述

有数字 $1\sim9$，每个数字的个数分别为 $cnt_1,cnt_2,cnt_3,...,cnt_9$ 。计算出“满意的集合“的个数。

"满意的集合" 定义：选出的数存在一种排列方式，其拼接起来后表示的十进制整数，能被 $3$ 整除，例如集合 $\{3,3,6\}$ $($包含了 $2$ 个数字 $3,1$ 个数字 $6$ $)$，可以有排列 $\{6,3,3\}$ 代表十进制下的整数 $633$，能被 $3$ 整除。

两个集合相同，当且仅当集合元素个数相同，且排序后对应数字相同，例如 $\{3,3,6\}$ 和 $\{3,6,3\}$ 是同样的集合。

空集合看作 $0$ ，是合法的，答案对 $1e9+7$ 取模。

输入描述:

输入一行，包括 $9$ 个整数 $cnt_1,cnt_2,cnt_3,...,cnt_9$，分别表示数字 $1\sim9$ 的个数，$0\le cnt_i\le 1e9$。

输出描述:

输出一行，表示”满意的集合”的个数，答案对 $1e9+7$ 取模。

示例1

输入

1 1 1 0 0 0 0 0 0

输出

4

说明

全部的数字集合 $\{1,2,3\}$。可能的集合有$\{\},\{1\},\{2\},\{3\},\{1,2\},\{1,3\},\{2,3\},\{1,2,3\}$。其中“满意的集合“有 $\{1,2\},\{3\},\{1,2,3\},\{\}$ 共 $4$ 个。

我的笔记

分析

首先简化题意，将“排列 $\{6,3,3\}$ 代表十进制下的整数 $633$，能被 $3$ 整除”，转化为“$\{6,3,3\}$ 中元素每一位的和能被 $3$ 整除”。

然后进行分析，本题可使用动态规划算法，将其转化为递推。递推第 0 层就是数字集合里不选任何数的情况，第 1 层就是数字集合里只选 1 的情况，第 2 层就是数字集合里只选 1、2 的情况，第 3 层就是数字集合里只选 1、2、3 的情况……（每层递推都允许不选任何数）

然后得到以下 DP 写法：

// TLE Version DP
for (int i = 1; i <= 9; i++) // 递推层数，直到推到第9层，即获得符合题意的答案
{
    for (int j = 0; j <= cnt[i]; j++) // 遍历选0个i、1个i、2个i……直到全选的所有情况
    {
        for (int k = 0; k < 3; k++) // 递推模为0、为1、为2的结果
        {
            dp[i][(j * i + k) % 3] += dp[i - 1][k];
            dp[i][(j * i + k) % 3] %= MOD; // MOD为1e9+7
        }
    }
}

但可以发现，该方法时间复杂度极高，为 $O(\sum_{n=1}^9cnt_n)$ ，一定会 TLE

进一步分析，每个数选 $k,k+3,k+6,\cdots,k+3*x(k\in\{0,1,2\})$ 时，和与 3 的模数是相等的。如：

(0 * 2) % 3 = 0, (1 * 2) % 3 = 2, (2 * 2) % 3 = 1,
(3 * 2) % 3 = 0, (4 * 2) % 3 = 2, (5 * 2) % 3 = 1……

这样我们就能把 $cnt_1$ 种情况转化为 3 种情况，即选 0 个、选 1 个、选 2 个，然后剩余情况都与这三个相等，我们统计一下这三种情况的个数，修改掉上面的 DP 中的第二层循环，就能将时间复杂度降低至 $O(9*3*3)$

源码

#include <bits/stdc++.h>
#define MOD 1000000007

using namespace std;

int main(void)
{
    int cnt[10];
    for (int i = 1; i <= 9; i++)
    {
        cin >> cnt[i];
    }
    long long dp[10][3];
    memset(dp, 0, sizeof(dp));
    dp[0][0] = 1;
    for (int i = 1; i <= 9; i++)
    {
        int mod_1 = (0 * i) % 3, mod_2 = (1 * i) % 3, mod_3 = (2 * i) % 3;
        int cnt_mod_1 = cnt[i] / 3 + 1, cnt_mod_2 = (cnt[i] + 1) / 3, cnt_mod_3 = (cnt[i] + 2) / 3;
        for (int j = 0; j < 3; j++)
        {
            dp[i][(j + mod_1) % 3] += dp[i - 1][j] * cnt_mod_1;
            dp[i][(j + mod_2) % 3] += dp[i - 1][j] * cnt_mod_2;
            dp[i][(j + mod_3) % 3] += dp[i - 1][j] * cnt_mod_3;
            dp[i][(j + mod_1) % 3] %= MOD;
            dp[i][(j + mod_2) % 3] %= MOD;
            dp[i][(j + mod_3) % 3] %= MOD;
        }
    }
    cout << dp[9][0] << endl;
    return 0;
}