Skip to content

2075C. Two Colors

binary search, combinatorics, math, 1500, https://codeforces.com/contest/2075/problem/C

Monocarp has installed a new fence at his summer house. The fence consists of 𝑛 planks of the same size arranged in a row.

Monocarp decided that he would paint his fence according to the following rules:

  • each plank of the fence will be painted in exactly one color;
  • the number of different colors that the planks will be painted in is exactly two;
  • the planks of the fence that are painted in the same color must form a continuous sequence, meaning that for all pairs of planks painted in the same color, there will be no planks painted in a different color between them.

Monocarp has 𝑚 different paints, and the paint of the 𝑖-th color is sufficient to paint no more than 𝑎𝑖 planks of the fence. Monocarp will not buy any additional paints.

Your task is to determine the number of different ways to paint the fence that satisfy all of Monocarp's described wishes. Two ways to paint are considered different if there exists a plank that is painted in different colors in these two ways.

Input

The first line contains a single integer 𝑡 (1≤𝑡≤10^4) — the number of test cases.

The first line of each test case contains two integers 𝑛 and 𝑚 (2≤𝑛,𝑚≤2⋅10^5) — the number of planks in the fence and the number of different colors of paint that Monocarp has.

The second line contains 𝑚 integers 𝑎1,𝑎2,…,𝑎𝑚 (1≤𝑎𝑖≤𝑛), where 𝑎𝑖 is the maximum number of planks that can be painted with the paint of color 𝑖.

The sum of 𝑛 over all test cases does not exceed 2⋅10^5. The sum of 𝑚 over all test cases does not exceed 2⋅10^5.

Output

For each test case, output the number of different ways to paint the fence that satisfy all of Monocarp's described wishes.

Example

Input

3
5 2
2 4
5 2
3 4
12 3
5 9 8

Output

4
6
22

Note

In the first test case, there are 4 different ways to paint the fence (the sequences of color numbers in which the planks can be painted from left to right are listed below):

  1. [1,2,2,2,2];
  2. [1,1,2,2,2];
  3. [2,2,2,1,1];
  4. [2,2,2,2,1].

In the second test case, there are 6 different ways to paint the fence (the sequences of color numbers in which the planks can be painted from left to right are listed below):

  1. [1,2,2,2,2];
  2. [1,1,2,2,2];
  3. [1,1,1,2,2];
  4. [2,2,1,1,1];
  5. [2,2,2,1,1];
  6. [2,2,2,2,1].

使用m-bisect_left(a,k)找到能涂k个板子的颜色种类数目,然后初步数目为x*y,假设k>n-k,那么能涂k块的x种颜色一定也可以涂(n-k)块,也就是说x中包含了y,所以要减去min(x,y),这些是用同一种颜色涂的方案数。

python
from bisect import bisect_left

for _ in range(int(input())):
    n, m = map(int, input().split())
    a = sorted(list(map(int, input().split())))
    ans = 0
    for k in range(1, n):
        x = m - bisect_left(a, k)
        y = m - bisect_left(a, n - k)
        ans += x * y - min(x, y)
    print(ans)

这段代码核心思路是从分割点 k 的角度出发,而不是从颜色对的角度出发。

以下是对这段代码逻辑的详细解读:

  1. 核心逻辑:枚举分割点 k

围栏由 n 块木板组成。由于必须使用恰好两种颜色且每种颜色必须连续,那么围栏一定被分成了两部分:

  • 左边部分:由颜色 C1k 块木板。
  • 右边部分:由颜色 C2nk 块木板。

这里 k 的取值范围是 [1,n1](保证左右两部分都至少有一块木板,从而保证使用了两种颜色)。

  1. 计算可用的颜色数量

对于一个确定的分割点 k

  • 能涂左边(k 块)的候选颜色数量为 x。在排好序的数组 a 中,满足 aik 的颜色数量即为 m - bisect_left(a, k)
  • 能涂右边(nk 块)的候选颜色数量为 y。同理,满足 ajnk 的颜色数量为 m - bisect_left(a, n - k)
  1. 排除“同一种颜色”的情况

根据乘法原理,如果我们任选一种符合左边的颜色 i 和一种符合右边的颜色 j,总共有 x×y 种组合。

但题目要求必须使用两种不同的颜色,因此必须排除 i=j 的情况。

  • 如果一种颜色 i 既能涂 k 块,又能涂 nk 块,那么它在 x×y 的组合中就被算成了一次“左右颜色相同”的情况。
  • 满足 aikaink 的颜色,实际上就是满足 aimax(k,nk) 的颜色。
  • 在已排序的数组中,满足 aik 的有 x 个,满足 aink 的有 y 个。那么同时满足这两个条件的颜色数量正好就是 min(x,y)

所以,对于确定的 k,合法的颜色对数量为:

count=x×ymin(x,y)
  1. 复杂度分析
  • 时间复杂度:对于每个测试用例,遍历 k1n1,循环内使用二分查找 bisect_left
    • 排序:O(MlogM)
    • 循环:O(NlogM)
    • 总计:O((N+M)logM)。给定 N,M2105,这个复杂度在 Python 中是可以过关的(约为 2105×183.6106 次操作)。
  • 空间复杂度O(M),用于存储颜色数组。
  1. 代码优化建议

代码逻辑很清晰,但在 Python 中,如果 N 很大,频繁调用 bisect_left 可能会稍微有点慢。由于 k 是线性递增的,nk 是线性递减的,其实可以使用双指针来替代二分查找,将复杂度降到 O(N+MlogM)

不过,就本题 1500 的难度和时限而言,这段 O(NlogM) 代码已经非常优秀且易于维护。

总结:这段代码巧妙地利用了“补集思想”(总组合减去相同颜色组合),将复杂的计数问题转化为了简单的统计问题。

为了使用双指针进一步优化,我们可以观察到随着分割点 k 的增加:

  1. k 是递增的:这意味着满足 aik 的最小索引 ptr1向右移动(或不动)。
  2. nk 是递减的:这意味着满足 ajnk 的最小索引 ptr2向左移动(或不动)。

通过维护这两个指针,我们可以将每个测试用例中循环部分的复杂度从 O(NlogM) 降低到 O(N+M)

双指针优化后的 Python 代码

python
import sys

def solve():
    # 使用快速读取
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    it = iter(input_data)
    t = int(next(it))
    
    output = []
    for _ in range(t):
        n = int(next(it))
        m = int(next(it))
        # 排序是必须的,O(M log M)
        a = sorted([int(next(it)) for _ in range(m)])
        
        ans = 0
        # ptr1 指向第一个满足 a[ptr1] >= k 的位置
        ptr1 = 0
        # ptr2 指向第一个满足 a[ptr2] >= n - k 的位置
        # 随着 k 增大,n-k 减小,ptr2 会向左移动,所以初始设在最右侧
        ptr2 = m 
        
        # 遍历分割点 k 从 1 到 n-1,O(N + M)
        for k in range(1, n):
            target_left = k
            target_right = n - k
            
            # 更新 ptr1:k 增大,ptr1 只能右移
            while ptr1 < m and a[ptr1] < target_left:
                ptr1 += 1
            
            # 更新 ptr2:n-k 减小,ptr2 只能左移
            while ptr2 > 0 and a[ptr2 - 1] >= target_right:
                ptr2 -= 1
            
            # x: 满足 a[i] >= k 的颜色数量
            x = m - ptr1
            # y: 满足 a[j] >= n-k 的颜色数量
            y = m - ptr2
            
            # 组合数 = x*y - 相同颜色的情况
            # 相同颜色的数量即为同时满足 >= k 和 >= n-k 的颜色数,即 min(x, y)
            ans += x * y - min(x, y)
            
        output.append(str(ans))
    
    sys.stdout.write("\n".join(output) + "\n")

if __name__ == "__main__":
    solve()

优化点解析

  1. 复杂度提升

    • 二分代码:每个 k 都进行一次 O(logM) 的二分查找,总复杂度 O(NlogM)
    • 新代码ptr1 在整个 k 循环中最多从 0 移动到 mptr2 最多从 m 移动到 0。因此循环部分的摊还复杂度是 O(N+M)
    • 总复杂度:从 O(T(MlogM+NlogM)) 优化到了 O(T(MlogM+N))
  2. 指针逻辑

    • ptr1 逻辑:当 k 变大时,之前满足条件的某些 ai 可能不再满足 aik,所以指针向后寻找。
    • ptr2 逻辑:当 k 变大时,nk 变小,原本不满足条件的某些 aj 现在可能满足 ajnk 了,所以指针向前寻找。
  3. Python 运行效率: 在 Codeforces 等平台,Python 的循环开销较大。使用双指针减少了 bisect_left 的函数调用开销,通常能获得更稳健的运行时间,尤其是在 N 很大时。使用 sys.stdin.read().split()sys.stdout.write 也是处理大数据量(2105)的常用技巧。