美团240309春招实习笔试——小美的区间删除

题目描述

小美拿到了一个大小为 n 的数组，她希望删除一个区间后，使得剩余所有元素的乘积末尾至少有 k 个 0。小美想知道，一共有多少种不同的删除方案？

输入描述

第一行输入两个正整数 n 和 k。

第二行输入 n 个正整数 $a_i$，代表小美拿到的数组。

约束条件

$1 ≤ n, k ≤ 10^5$
$1 ≤ a_i ≤ 10^9$

输出描述

一个整数，代表删除的方案数。

示例 1

输入

5 2 
2 5 3 4 20

输出

4

说明

第一个方案，删除 [3]。

第二个方案，删除 [4]。

第三个方案，删除 [3,4]。

第四个方案，删除 [2]。

解题

末尾零的个数由所有乘数中的2因子和5因子决定，公式化来说 $$tailing_zero_num = min(\sum_0^n factor2(num[i]), \sum_0^n factor5(nums[i]))$$
设原数组所有2因子和5因子的和分别为total2和total5，被删除区间中2因子和5因子的和分别为x2和x5。
根据题意，要使末尾零不少于k，等价于$min(total2-x2, total5-x5) >= k$
即$total2-x2 >= k$ and $total5-x5 >= k$, 也即$x2 <= total2 - k$ and $x5 <= total5 - k$
换句话说，原题等价于寻找2因子和5因子满足上式的子区间。

首先计算出相应值：

from bisect import bisect_right
from itertools import accumulate

n, k = 5, 2
nums = [2, 5, 3, 4, 20]

def factor_nums(n: int, factor: int):
    num = 0
    while n and not n % factor:
        num += 1
        n //= factor
    return num

factor_2_nums = [factor_nums(num, 2) for num in nums]
factor_5_nums = [factor_nums(num, 5) for num in nums]

total2, total5 = sum(factor_2_nums), sum(factor_5_nums)

首先尝试使用滑动窗口：

ans = 0
i, j = -1, -1
x2, x5 = 0, 0
while True:
    print(i, j)
    if x2 <= total2 - k and x5 <= total5 - k:
        if j > i:
            ans += 1
            print(i, j, nums[i+1: j+1])
        j += 1
        if j == len(nums): break
        x2 += factor_2_nums[j]
        x5 += factor_5_nums[j]
    else:
        i += 1
        x2 -= factor_2_nums[i]
        x5 -= factor_5_nums[i]

print(ans)

后发现错误，如当前区间为[3,4]时，满足条件，下一个待验证区间为[3,4,5]，答案区间[4]被略过。

那么直接遍历所有子区间：

ans = 0
for i in range(n):
    x2, x5 = 0, 0
    for j in range(i, n):
        x2 += factor_2_nums[j]
        x5 += factor_5_nums[j]
        if x2 <= total2 - k and x5 <= total5 - k:
            print(i, j, nums[i: j+1])
            ans += 1
        else:
            break
    
print(ans)

这样时间复杂度为 $O(n^2)$

改进：使用前缀和+二分查找

prefix2 = list(accumulate(factor_2_nums, initial=0))
prefix5 = list(accumulate(factor_5_nums, initial=0))
ans = 0
for i in range(n+1):  # 固定左端点，二分查找右端点
 	# bisect_right找到的是满足大于条件的第一个下标.第二个参数可以理解为寻找从prefix2[i]开始，与其相差total2-k的元素
    j2 = bisect_right(prefix2, total2 - k + prefix2[i]) 
    j5 = bisect_right(prefix5, total5 - k + prefix5[i])
    ans += min(j2, j5) - i - 1
    print(nums[i: min(j2, j5) - 1])
print(ans)

时间复杂度为$O(nlogn)$

参考博客