黑白气球

1 教学目标

• 公式$\sum s(i)$,$s(i)$:以$i$为结尾的合法对数
• 本质不重不漏集合分类问题
• 固定结尾,用 map 作为桶维护前面每种颜色的数量
• 把逐个结尾的贡献合并成组合数,理解 $\binom{c}{2}$ 的来源

2 题目: 黑白气球

有一排黑白气球,你可以随机选两个球,问有多少种方法选到颜色不同的球?

• 第一行表示有多个气球
• 0表示白色
• 1表示黑色

1
2

6
0 1 1 0 0 1

输出

1

9

代码

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

//Author by [Rainboy](https://github.com/rainboylvx) 
//date: 2024-05-03 17:22:44
#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e6+5;
int n,m;
int a[maxn];
int ans;

int main (int argc, char *argv[]) {
    std::cin >> n;
    for(int i = 1;i <= n ;++i ) // i: 1->n
    {
        std::cin >> a[i];
    }

    //枚举以i为结尾
    for(int i = 2;i <= n ;++i ) // i: 1->n
    {
        // 枚举对应的起点
        for(int j = 1;j <= i-1 ;++j ) // j: 1->i
        {
            if( a[j] != a[i])
                ans++;
        }
    }
    std::cout << ans << "\n";

    return 0;
}

3 解释

本质这是一个不重不漏的集合分类问题

先给每个气球进行编号

1
2

0 1 1 0 0 1
1 2 3 4 5 6

那么答案的$9$对气球如下:

3.1 证明这是一个集合问题:

原题目部我们: 1. 任何选一对 2. 且这一对的颜色不同

我们一步一步做,任选一对气球,可以得到集合$A = \{(1,2),{1,3},{2,3},{1,4},{2,4}, \cdots,(5,6)\}$这样的集合,然后去除集合$A$中颜色相同的一对气球,得到集合$B = \{(1,2), (1,3), (2,4), (3,4), (2,5), (3,5), (1,6), (4,6), (5,6) \}$,那么问题的答案就是集合$B$的基数(元素数量)

3.2 集合分类

如果我们可以把集合$B$不重不漏的分成多个子集合$b_i$,那么答案就是$\sum |b_i|$

显然按对数的结尾的数字,可以把集合$B$不重不漏的分成多个子集$b_i$.$b_i$表示结尾为$i$的元素形成的集合.例如$b_3 = \{(1,3)\}$.

证明:

• 不重: 显然以$i$为结尾的元素不可能和以$j$为结尾的元素在同一个集合
• 不漏: 任意一个元素一定会以$i \in [1,n]$为结尾

4 对数问题的公式: $\sum s(i)$

把上面的做法抽象出来.设 $s(i)$ 表示以 $i$ 为结尾的合法对数,也就是

那么答案就是把每个结尾位置的贡献加起来:

这就是对数问题的公式 $\sum s(i)$.

4.1 公式为什么是对的

• 不漏: 任意一个合法数对 $(j,i)$ 一定有一个右端点 $i$,所以它一定会被某个 $s(i)$ 统计到
• 不重: 数对 $(j,i)$ 的右端点 $i$ 是唯一的,所以它只会被 $s(i)$ 统计一次,不会在别的子集里再算一遍
• $i=1$ 时左边没有任何元素,所以 $s(1)=0$;从 $i=2$ 开始累加结果也一样

这正是前面说的按结尾分类加上分类加法原理: 把集合 $B$ 不重不漏地切成 $b_1,b_2,\cdots,b_n$,答案就是 $\sum |b_i|$,而 $|b_i|$ 就是这里的 $s(i)$.

4.2 难点不在求和,而在怎么求 $s(i)$

公式本身只是"把每个结尾的答案加起来",真正的难点是如何快速算出一个 $s(i)$.

最直接的办法是再枚举一遍 $j$,那就是 $O(n^2)$:

1
2
3
4
5

// 暴力:以 i 为结尾,再枚举起点 j
for (int i = 2; i <= n; i++)
    for (int j = 1; j < i; j++)
        if (a[j] != a[i])
            ans++;          // 这里就是在算 s(i)

5 方法一: 固定结尾,用 map 作为桶

固定第 $i$ 个气球作为结尾,需要统计前面有多少个气球能和它配对.暴力做法每次都重新看一遍前面的气球,但我们只关心各颜色的数量,可以把这些数量保存下来.

用 map<int, long long> cnt 作为桶:cnt[color] 表示当前位置之前,颜色为 color 的气球数量.这里的键表示颜色,值表示数量;没有出现过的颜色,通过 cnt[color] 访问时会得到初始值 $0$.

处理第 $i$ 个气球之前,前面一共有 $i-1$ 个气球,其中同色的有 $cnt[a_i]$ 个,所以能够配对的数量是

算完贡献以后,再把当前气球放进桶里,供后面的气球使用.

1
2
3
4
5
6

map<int, long long> cnt;
long long ans = 0;
for (int i = 1; i <= n; i++) {
    ans += (i - 1) - cnt[a[i]];  // 只统计前面与当前气球异色的气球
    cnt[a[i]]++;                 // 当前气球加入桶,供下一个结尾使用
}

对题目中的例子,整个计数过程如下.表中的桶数量都是加入当前气球之前的数量.

把最后一列相加,得到 $0+1+1+2+2+3=9$.

先统计、再入桶,保证桶里始终只有前面的气球.如果把两行交换,当前气球会被当成前面的同色气球减掉,每个结尾都会少算 $1$.

如果共有 $k$ 种颜色,std::map 每次访问需要 $O(\log(k+1))$,总时间为 $O(n\log(k+1))$,空间为 $O(k)$.本题只有黑白两种颜色,也可以直接用 long long cnt[2] = {}; 作为桶,总时间为 $O(n)$.

桶里应当保存什么,取决于当前气球需要什么样的配对对象:

6 方法二: 合并结尾的贡献,用组合数计数

继续从固定第 $i$ 个气球作为结尾出发:

把所有结尾的贡献加起来,可以拆成两部分:

这里的 $cnt[a_i]$ 仍然表示处理位置 $i$ 之前的同色气球数量.第一部分统计所有气球对,第二部分统计同色气球对.

先看所有气球对.每个结尾前面的气球数依次为 $0,1,2,\ldots,n-1$,因此

再看某一种颜色.假设这种颜色共有 $c$ 个气球,按它们在原序列中出现的先后顺序,前面的同色气球数依次为 $0,1,2,\ldots,c-1$.即使中间夹着其他颜色,这个规律也成立,所以这种颜色产生的同色对数为

这也可以理解为: 从这 $c$ 个气球中任选两个,位置靠后的那个就是结尾,每一对恰好被统计一次.

固定一个结尾时,我们只需从前面选一个气球;把一整组气球的所有结尾贡献加起来,才得到从这一组中选两个的组合数.

设白球共有 $w$ 个,黑球共有 $b$ 个,$w+b=n$.所有气球对去掉两种同色对,剩下的就是异色对:

样例中 $n=6,w=b=3$,所以答案为 $15-3-3=9$.

1
2
3
4
5
6
7
8

long long cnt[2] = {};
for (int i = 1; i <= n; i++)
    cnt[a[i]]++;

long long total = n;
long long ans = total * (total - 1) / 2;
for (int color = 0; color < 2; color++)
    ans -= cnt[color] * (cnt[color] - 1) / 2;

读入并统计颜色需要 $O(n)$ 时间,套公式只需 $O(1)$ 时间,额外空间为 $O(1)$.数量和答案使用 long long,避免乘法溢出.

把公式展开,还能得到 $ans=wb$: 每个白球都能和每个黑球配对.如果有多种颜色,则仍然可以使用“所有对数减去每种颜色的同色对数”的公式.

两种方法计算的是同一个答案: 用桶时,每遇到一个结尾就累加贡献;用组合数时,先统计每种颜色的总数,再把这些贡献一次合并.

7 总结

• 对数类问题的通常解法: 按结尾位置把集合分类,答案就是 $\sum s(i)$
• 关键在于如何快速求出一个子集的数量 $s(i)$(通常用桶)
• 同一组内任意两个元素都能配对时,可以把各个结尾的贡献合并为 $\binom{c}{2}$;本题用总对数减去同色对数
• 这里用到了集合与分类加法原理

8 练习题目

建议按下面的顺序练习.先把本文的两种计数方法写熟,再练习从配对条件中找出桶的键.

8.1 巩固: 从颜色分类到组合数

1. atcoder abc206_c: [ABC206C] Swappable

给定一个整数序列,统计 $i<j$ 且 $a_i\ne a_j$ 的数对数量.这相当于把黑白气球推广到多种颜色,是最适合接着本文练习的一题.

尝试分别用“固定结尾、用桶统计”和“总对数减去同值对数”写两份程序.数值可达 $10^9$,可以用 map 记录次数;答案需要 long long.

2. atcoder abc159_d: [ABC159D] Banned K

对每个位置 $k$,分别计算去掉第 $k$ 个球后,剩余球中数值相同的球对有多少.每次删除都从原序列出发,各次询问互不影响.

训练点是组合数的局部变化: 先用 $\sum_x\binom{cnt[x]}{2}$ 算出原来的同值对数.去掉一个球时,只有它所属的那一组会受影响,想一想有多少对包含这个球.

8.2 迁移: 根据配对条件查桶

3. luogu P1102: A-B 数对（洛谷原题）

统计满足 $A-B=C$ 的数对,其中 $C>0$,相同数值出现在不同位置时要分别计数.本题没有要求两个数按输入位置满足 $i<j$.

固定一个 $A$,需要找的数值就是 $A-C$.可以先统计整个序列中每个数值的出现次数,再逐个枚举 $A$ 查询桶.注意桶里存的是次数,只判断某个数是否存在会漏掉重复元素产生的配对.

4. codeforces 1520D: Same Differences

统计 $i<j$ 且 $a_j-a_i=j-i$ 的数对数量.先移项,把条件变成 $a_i-i=a_j-j$,再把 $a_i-i$ 当作一种新的“颜色”.

尝试再次用本文的两种方法求解: 扫描时查询前面同类元素的数量,或者统计每一类的大小后累加组合数.这道题练习的是先变换条件,再分类计数.

8.3 进阶: 变换后的键与复合键

5. atcoder abc166_e: [ABC166E] This Message Will Self-Destruct in 5s

选择两个人,要求编号之差的绝对值等于两人的身高之和.固定编号较大的 $i$,前面的 $j$ 需要满足 $i-j=a_i+a_j$.

把条件整理成 $j+a_j=i-a_i$: 扫到 $i$ 时查询桶中的 $i-a_i$,然后把 $i+a_i$ 放入桶.这里查询的键和存入的键不同,要先明确两者各自代表什么.

6. luogu P3405: [USACO16DEC] Cities and States S（USACO 原题）

两个城市的名称前两个字母,分别等于对方的州代码,且两个城市必须来自不同州.把每个城市表示成“名称前缀、州代码”这一对信息.

用 map<pair<string, string>, long long> 作为桶: 当前信息为 $(p,s)$ 时,查询前面有多少个 $(s,p)$,再把当前信息入桶.若 $p=s$,它无法参与合法配对,应跳过.这道题把一个数作为键推广为一对字符串作为键,同时练习先查询、后入桶来避免重复计数.

做每道题时,都先写清楚合法条件、固定一个对象后的贡献,以及桶的键和值分别表示什么.再判断这些贡献能否合并成组合数.

9 相关题解