3524. 求出数组的 X 值 I

发表于 分类于 阅读次数： Waline： 本文字数： 919 阅读时长 ≈ 3 分钟

考点

  • 子数组之间的状态转移
  • 乘法取模

思路

一、问题重述与等价转化

给定一个整数数组 \[ \text{nums} = [a_1, a_2, \dots, a_n] \] 以及一个正整数 \(k\)。

允许从数组中删除一个前缀和一个后缀（均可为空），剩下的部分必须是一个非空连续子数组。 对所有可能得到的子数组，统计其元素乘积对 \(k\) 取模后的结果分布。

形式化地，对所有子数组 \([l, r]\)（\(1 \le l \le r \le n\)），定义 \[ P(l,r) = \prod_{i=l}^r a_i \] 需要统计每个 \[ x \in \{0,1,\dots,k-1\} \] 满足 \[ P(l,r) \bmod k = x \] 的子数组数量。

二、核心数学性质

整个问题仅依赖以下基本模运算性质： \[ (a \cdot b) \bmod k = \big( (a \bmod k)\cdot(b \bmod k) \big) \bmod k \] 因此，在计算子数组乘积时：

  • 不需要关心真实乘积的大小
  • 只需要维护“乘积对 \(k\) 的余数”

三、问题建模：按右端点分类的动态规划

3.1 状态定义

设在处理到数组第 \(i\) 个元素时（即右端点为 \(i\)）： \[ \text{dp}_i[r] = \text{以 } i \text{ 结尾的子数组中，乘积对 } k \text{ 取模为 } r \text{ 的个数} \] 其中： \[ r \in \{0,1,\dots,k-1\} \] 最终答案数组为： \[ \text{ans}[r] = \sum_{i=1}^n \text{dp}_i[r] \]

3.2 状态转移的来源

考虑第 \(i\) 个元素 \(a_i\)，记： \[ x = a_i \bmod k \] 以 \(i\) 为右端点的子数组来源只有两类：

情况一：新开子数组

只包含当前元素的子数组： \[ [i,i] \] 其乘积余数为： \[ x \] 贡献为： \[ \text{dp}_i[x] \mathrel{+}= 1 \]

情况二：扩展已有子数组

对所有以 \(i-1\) 结尾、乘积余数为 \(j\) 的子数组： \[ \text{dp}_{i-1}[j] \] 将其右端点扩展到 \(i\)，新余数为： \[ (j \cdot x) \bmod k \] 转移方程： \[ \text{dp}_i[(j\cdot x)\bmod k] \mathrel{+}= \text{dp}_{i-1}[j] \]

3.3 完整状态转移方程

综合两类情况，可得： \[ \boxed{ \begin{aligned} \text{dp}_i[r] &= \mathbf{1}_{r=x} + \sum_{j=0}^{k-1} \Big[ (j\cdot x)\bmod k = r \Big] \cdot \text{dp}_{i-1}[j] \end{aligned} } \] 其中 \(\mathbf{1}_{r=x}\) 表示当 \(r=x\) 时取 1，否则取 0。

四、算法流程

  1. 使用长度为 \(k\) 的数组 pre_cnt 表示 \(\text{dp}_{i-1}\)
  2. 使用长度为 \(k\) 的数组 cur_cnt 表示 \(\text{dp}_i\)
  3. 对数组从左到右遍历：
    • 清空 cur_cnt
    • 加入新子数组 [i,i]
    • 扩展所有旧子数组
    • 将 cur_cnt 累加到答案数组
    • 交换 pre_cnt 与 cur_cnt

五、复杂度分析

  • 状态数：\(k \le 5\)
  • 每个元素进行一次 \(k\) 规模转移

时间复杂度： \[ O(n \cdot k) \] 空间复杂度： \[ O(k) \]

六、AC代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class Solution {
public:
    vector<long long> resultArray(vector<int>& nums, int k) {
        // ans[r]: 所有子数组中，乘积 % k == r 的个数
        vector<long long> ans(k, 0);

        int n = nums.size();

        // 使用固定大小数组，避免 vector 带来的动态分配开销
        // pre_cnt[r]: 以上一位置为右端点、乘积 % k == r 的子数组个数
        // cur_cnt[r]: 当前处理位置为右端点的对应统计
        long long arr0[6] = {}, arr1[6] = {};
        long long *pre_cnt = arr0, *cur_cnt = arr1;

        // 枚举右端点 i
        for (int i = 1; i <= n; ++i) {
            int x = nums[i - 1] % k;

            // 清空当前轮 DP 状态
            memset(cur_cnt, 0, sizeof(long long) * 6);

            // 情况一：新开子数组 [i, i]
            cur_cnt[x] += 1;

            // 情况二：扩展所有以 i-1 结尾的子数组
            for (int j = 0; j < k; ++j) {
                if (pre_cnt[j] == 0) continue;
                int r = (int)(1LL * j * x % k);
                cur_cnt[r] += pre_cnt[j];
            }

            // 当前轮所有子数组均为合法答案，累加到总统计中
            for (int j = 0; j < k; ++j) {
                ans[j] += cur_cnt[j];
            }

            // 滚动数组，准备处理下一个位置
            swap(pre_cnt, cur_cnt);
        }

        return ans;
    }
};