蓝桥杯真题解题思路——因数计数

前言

  印象里这题我在 24 年 4 月的蓝桥杯 C++ 本科生 A 组见到过，但是当时我实在太菜了，在赛场上没有想出来。时隔一年我重温这道题，没有借鉴任何人的思路，一上来就奔着 $O(n\log n)$ 的复杂度，虽然最后实际上是 $O(n{\log }^{2}n)$ 的复杂度，但也算并独立完成了这道题的求解。这个时间复杂度可能不是最优的，但是应该是比较好理解的。

思路

  首先，我们可以理解一下题目的意思。这里是题目链接。

求解二元组的个数

  首先，题目里提到了一个所谓的“二元组”。

小蓝随手写出了含有 $n$ 个正整数的数组 { a 1 , a 2 , ⋯   , a n } \{a_1,a_2,\cdots,a_n\} {a1​,a2​,⋯,an​}，他发现可以轻松地算出有多少个有序二元组 $(i,j)$ 满足 $a_j$ 是 $a_i$ 的一个因数。

  我们肯定要先研究这些二元组，再去研究后面四元组的事情。求解二元组说起来容易，可能是因为小蓝写的数组太小了，$O(n^2)$ 的时间复杂度也无所谓。但是限制于本题的数据，我们最好是能找到一种 $O(n\log n)$ 的方式求出这些二元组。
  那么这样的二元组有多少个呢？我们可以定义一些符号：

• $p_i$：满足 $a_j(j\ne i)$ 是 $a_i$ 的倍数的 $j$ 的个数。即数组里面除了 $a_i$，其它这 $n-1$ 个里面有多少个为 $a_i$ 的倍数。
• $q_i$：满足 $a_j(j\ne i)$ 是 $a_i$ 的因数的 $j$ 的个数。即数组里面除了 $a_i$，其它这 $n-1$ 个里面有多少个为 $a_i$ 的因数。

  显然，这些二元组的总数应为 $s={\sum }_{i=1}^n{p}_i={\sum }_{i=1}^n{q}_i$。

为什么 ${\sum }_{i=1}^n{p}_i={\sum }_{i=1}^n{q}_i$？你可以理解为，当我们找到一个二元组 $(i,j)$ 时，$p_i$ 将增加 $1$，而 $q_j$ 也会增加 $1$。每发现一个二元组，都会导致 $\sum p$ 和 $\sum q$ 都增加 $1$，所以有 ${\sum }_{i=1}^n{p}_i={\sum }_{i=1}^n{q}_i$。

  那么怎么求解 ${\sum }_{i=1}^n{p}_i$ 或者 ${\sum }_{i=1}^n{q}_i$？我们可以先对数组 $\mathbf{a}$ 从小到大排序，这对答案显然不会有影响，排序后的数组仍然记为 $\mathbf{a}$。
  接下来令 $i$ 从 $1$ 遍历到 $n$，分别去求解 $p_i$，然后全部求和就是 $s$。说到求解 $p_i$，一个朴素的想法是：遍历 $a_{i+1}$ 到 $a_n$，一个一个判断是不是 $a_i$ 的倍数。很遗憾，这种方法时间复杂度过高，为 $O(n^2)$。
  那我们不妨换一种思路。我要求 $p_i$，也就是 $a_i$ 的倍数的个数，我可以看数组中等于 $2a_i$ 的有多少个，再看等于 $3a_i$ 的有多少个……一直这样找下去，然后全部加起来就是 $p_i$。由于现在的 $\mathbf{a}$ 是一个升序数组，使用二分查找，可以使得找数组中等于 $t{a}_i$ 元素个数的时间复杂度降到 $O(\log n)$。这个 $t$ 是从 $2$ 开始枚举，然后是 $3$，那么一直枚举到多少可以结束呢？如果记 $\mathbf{a}$ 中最大的元素是 $m$，那么答案是 $\lfloor m/a_i\rfloor$。
  根据上面的思路，求解 $p_i$ 的时间复杂度应该是 $O((m\log n)/a_i)$。那么求解 $s={\sum }_{i=1}^n{p}_i$，时间复杂度应该就是：$O\left((m\log n){\sum }_{i=1}^n(1/a_i)\right)$。假设 $a_i$ 互不相等，那么通过基本的放缩，可以知道 ${\sum }_{i=1}^n(1/a_i)\le {\sum }_{i=1}^n(1/i)\sim O(\log n)$，因此求解 $s$ 的时间复杂度是 $O(m{\log }^{2}n)$。题目中 $m$ 和 $n$ 是同上限的，因此也可以将总时间复杂度写作 $O(n{\log }^{2}n)$。
  那如果 $a_i$ 有重复的，上面的放缩就不成立。但是重复的 $a_i$ 意味着其中某个的 $p_i$ 计算出来了，其它的直接照抄这个 $p_i$ 就行，此时就不需要额外算了。这样时间复杂度还是 $O(n{\log }^{2}n)$。
  顺带说一句这个 $q_i$，其实顺带就给算完了。比如说我们找到数组中等于 $2a_i$ 的元素为 $a_j$ 到 $a_k$。那么 $p_i$ 的变动是增加 $k-j+1$，而 $q_j$ 到 $q_k$ 都需要加上 $1$。由于每次对数组 $q$ 都是在某个连续区间上加上 $1$，因此可以构造 $q$ 的差分数组 $qdiff$，保证在常数时间复杂度内完成 $qdiff$ 的修改（否则会拖累时间复杂度到 $O(n^2)$）；所有 $a_i$ 遍历完了之后，使用 $O(n)$ 的时间从 $qdiff$ 构造 $q$ 即可。

求解四元组的个数

  然后题目问我们四元组 $(i,j,k,l)$ 有多少个。实际上我是这样想的，先确定了 $i$ 和 $k$，在此基础上想一下 $(j,l)$ 有多少种可能？答案是 $\pi (i,k)=s-p_i-q_i-p_k-q_k+1+\mathbb{I}(a_i=a_k)$。这里的 $\mathbb{I}(\cdot )$ 把一个 bool 值转换为整数，意思是当 $a_i=a_k$ 时，$\mathbb{I}(a_i=a_k)=1$，否则 $\mathbb{I}(a_i=a_k)=0$。
  为什么是这个结果？首先，可以确定 $a_k$ 是 $a_i$ 的倍数。如果我们随便填 $(j,l)$，显然有 $s$ 种填法。然而这其中有一些会导致四元组有重复元素，要抠掉。一共有四种情况需要抠掉：

• $j=i$：此时 $l$ 可以填的数字有 $p_i$ 种，所以这种情况有 $p_i$ 种。

• $j=k$：此时 $l$ 可以填的数字有 $p_k$ 种，所以这种情况有 $p_k$ 种。

• $j\ne i\wedge j\ne k\wedge l=k$：如果是单纯的 $l=k$，那么 $j$ 有 $q_k$ 种可能，并且这些填法保证 $j\ne k$。但是由于 $a_i$ 是 $a_k$ 的因数，所以这 $q_k$ 种可能包含 $j=i$ 的情况。那么这种情况就是 $q_k-1$ 种。

• $j\ne i\wedge j\ne k\wedge l=i$：如果是单纯的 $l=i$，那么 $j$ 有 $q_i$ 种可能，并且这些填法保证 $j\ne i$。

  • 如果 $a_k$ 也是 $a_i$ 的因数（考虑到 $a_i$ 是 $a_k$ 的因数，等价于 $a_k=a_i$），那么这 $q_i$ 种可能包含 $j=k$ 的情况。
  • 如果 $a_k$ 不是 $a_i$ 的因数，那么这 $q_i$ 种可能不包含 $j=k$ 的情况。

  因此，这种情况有 $q_i-\mathbb{I}(a_i=a_k)$ 种。

  那么回到我们的题目，最终四元组的个数就是遍历所有的二元组 $(i,k)$，并求和 $\pi (i,k)$。可以比表示为 $ans={\sum }_{(i,k)}\pi (i,k)={\sum }_{(i,k)}[s-p_i-q_i-p_k-q_k+1+\mathbb{I}(a_i=a_k)]$。
  我们可以把上面的和式拆开：

• ${\sum }_{(i,k)}(s+1)$：这也就是 $s(s+1)$。
• ${\sum }_{(i,k)}\mathbb{I}(a_i=a_k)$：数组 $\mathbf{a}$ 中值相同的元素分为一组。对于每一组而言，如果这一组有 $x$ 个元素，那么这一组的贡献就是 $x(x-1)$。由于同一组的元素在排序后的 $\mathbf{a}$ 中必连续，所以可以在 $O(n)$ 复杂度内解决。
• $-{\sum }_{(i,k)}(p_i+q_i+p_k+q_k)$：这就要求我们重新温故一下二元组的求取过程。之前是每找到一个二元组 $(i,k)$，执行 $p_i=p_i+1$ 和 $q_k=q_k+1$。现在我们每找到一个二元组 $(i,k)$，都需要计算一下 $-(p_i+q_i+p_k+q_k)$ 并加入答案。

  最终，我们可以在 $O(n{\log }^{2}n)$ 的时间复杂度内求得答案。

AC 代码

import os
import sys

def make_array(length,val):
    return [val for _ in range(length)]
def make_2d_array(rows,cols,val):
    return [[val for _ in range(cols)] for _ in range(rows)]
def make_3d_array(d1,d2,d3,val):
    return [[[val for _ in range(d3)] for _ in range(d2)] for _ in range(d1)]

def read_int():
    return int(input())
def read_ints():
    return [int(i) for i in input().split()]

from bisect import bisect_left,bisect_right
n = read_int()
a = sorted(read_ints())
m = max(a)
p = make_array(n,-1)
q = make_array(n,0)
sp = 0
qdiff = make_array(n,0)
qdiff[0] = -1
qsum = make_array(n + 1,0)
psum = make_array(n + 1,0)
ans = 0
for i in range(n):
    for j in range(1,m // a[i] + 1):
        target = a[i] * j
        lidx = bisect_left(a,target)
        if lidx != n and a[lidx] == target:
            ridx = bisect_right(a,target)
            p[i] += ridx - lidx
            qdiff[lidx] += 1
            if ridx < n:
                qdiff[ridx] -= 1
last = -1;count = 0
for i in range(n):
    if i == 0:
        q[0] = qdiff[0]
    else:
        q[i] = q[i - 1] + qdiff[i]
    sp += p[i]
    if a[i] != last:
        ans += count * (count - 1)
        last = a[i]
        count = 1
    else:
        count += 1
# print(p);print(q);print(qdiff);exit()
ans += sp * (sp + 1) + count * (count - 1)
# print(ans);exit()
for i in range(n):
    qsum[i + 1] = qsum[i] + q[i]
    psum[i + 1] = psum[i] + p[i]
for i in range(n):
    ans += 2 * (p[i] + q[i])
    for j in range(1,m // a[i] + 1):
        target = a[i] * j
        lidx = bisect_left(a,target)
        if lidx != n and a[lidx] == target:
            ridx = bisect_right(a,target)
            ans -= (p[i] + q[i]) * (ridx - lidx)
            ans -= psum[ridx] - psum[lidx] + qsum[ridx] - qsum[lidx]
print(ans)

  上面的代码存在一个漏洞，就是求 $s$ 的时候并没有考虑 $a_i$ 重复的情况，容易被卡。但是蓝桥杯还是比较仁慈的，没有这种恶心数据，最终还是过了。当然在比赛的时候，如果有时间，最好能够考虑这种情况，并且改进代码。