牛客小白月赛43

E - 满意的集合

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

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

题目描述

有数字 1∼9，每个数字的个数分别为 cnt1,cnt2,cnt3,...,cnt9 。计算出“满意的集合“的个数。

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

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

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

输入描述:

输入一行，包括 9 个整数 cnt1,cnt2,cnt3,...,cnt9，分别表示数字 1∼9 的个数，0≤cnti≤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(∑n=19cntn) ，一定会 TLE

进一步分析，每个数选 k,k+3,k+6,⋯,k+3∗x(k∈{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……

这样我们就能把 cnt1 种情况转化为 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;
}