小C选点 | MarsCode AI 刷题

问题描述

小C在坐标轴上有 n 个点，她想从中选出 m 个点，使得这些点之间的两两距离都不超过 k。你的任务是帮助小C计算出有多少种选点的方案。由于答案可能非常大，请将答案对 $10^9+7$ 取模。

例如：给定点的坐标为 [1, 2, 3, 4]，并且 k=3，你可以任选两个点，它们的距离都不会超过 k，因此总共有6种选点方案。

测试样例

样例1：

输入：n = 4 ,m = 2 ,k = 3 ,x = [1, 2, 3, 4]
输出：6

样例2：

输入：n = 4 ,m = 2 ,k = 2 ,x = [1, 2, 3, 4]
输出：5

样例3：

输入：n = 5 ,m = 3 ,k = 5 ,x = [1, 3, 6, 7, 9]
输出：3

分析

乍一看这道题目，尤其是因为10^9+7这个数字的出现，貌似计算量很大，如果没有设置好存储空间很容易out of memory。

那么有没有什么简化的办法呢？

当然有，那就是以空间换时间，我们知道如果从n个点中选取m个点可以直接使用组合数的公式$C(n,m)=\frac{m!}{(n−m)!n!}$。所以为了避免重复计算，可以预先计算出所有点的阶乘和取模。

def precompute_factorials(max_n):
    fact = [1] * (max_n + 1)
    for i in range(2, max_n + 1):
        fact[i] = fact[i - 1] * i % MOD
    return fact

但是对于模运算中直接计算除法是不可行的，所以在组合数计算中，我们需要用到分母的逆元来避免除法。

费马小定理：当 m 为质数且 a 不为 m 的倍数，有 $a^{m−1}≡1 mod (m)$ 另一个形式：对于任意整数 a ，有 $a^m≡a(modm)$ 。

根据费马小定理可知： $a^{m−2}$ 就是 a 在模 m 意义下的逆元。$a^{-1}=a^{m-2} mod(m)$。 使用pow可以高效地计算大数的幂和模。

def mod_inv(x, p):
    return pow(x, p - 2, p)

计算组合数：

def comb(n, m, fact):
    if m > n or m < 0:
        return 0
    return fact[n] * mod_inv(fact[m], MOD) % MOD * mod_inv(fact[n - m], MOD) % MOD

因为，根据题目条件，一次选取的殿中任意的两个点的最远距离不会超过k，我们可以将数组进行排序，复杂度$O(nlogn)$。之后再便利保证每次取点的区间段左右两端点距离不超过k，即可。（每次固定左端点，再在j-i-1区间选m-1个点）复杂度$O(n)$

直接放代码：

def solution(n: int, m: int, k: int, x: list) -> int:
    points = x.copy()
    n = len(points)
    points.sort()
    
    fact = precompute_factorials(n)
    total_ways = 0
    
    j = 0
    for i in range(n):
        while j < n and points[j] - points[i] <= k:
            j += 1
        # j is now one past the valid range
        total_ways += comb(j - i - 1, m - 1, fact)
        total_ways %= MOD
    
    return total_ways

尾言

小白作者，欢迎交流讨论