笔试 | 字节跳动2021秋招第三场 数数组题目详解&扩展
文章首发于公众号面鲸，欢迎关注。另外由面鲸组织的每天一道高频面试题刷题挑战活动正在进行中，快来挑战吧～
数数组
题目描述
- 现在要生成一个数列，满足以下条件
- 数列的长度为n，且每个数字的大小 a [ i ] a[i] a[i]满足 l < = a [ i ] < = r l<=a[i]<=r l<=a[i]<=r
- 数组的和要能被3整除
- 现在给定三个数 n , l , r n,l,r n,l,r，可否求出满足以上条件的数组个数，因为数组数量比较大，最终结果请mod 1e9+7
分析
- 遇到这种问题，一看是mod 1e9+7，那么结果一定很大，应该是有一定的递推式子的。
- 令 b 0 b_0 b0表示[l,r]中除3余0的数的个数， b 1 b_1 b1表示[l,r]中除3余1的数的个数， b 2 b_2 b2表示[l,r]中除3余2的个数；令 a r r [ i ] [ x ] arr[i][x] arr[i][x]表示前i位除3余x的方案数，有 a r r [ 1 ] [ x ] = b x arr[1][x]=b_x arr[1][x]=bx。一种最直接计算 b x b_x bx的方法是遍历[l,r]中的所有整数，挨个统计。但是注意到[l,r]的范围可能很大，因此不能直接这么做。那有什么办法可以快速计算得到这些数呢。如果我们要统计除3余1的数的个数，那么这个数可以表示为 3 k + 1 3k+1 3k+1，进而有 l < = 3 k + 1 < = r l<=3k+1<=r l<=3k+1<=r，也就是 ( l − 1 ) / 3 < = k < = ( r − 1 ) / 3 (l-1)/3 <=k<=(r-1)/3 (l−1)/3<=k<=(r−1)/3，又因为k必须是整数，所以k的个数为 ( r − 1 ) / 3 (r-1)/3 (r−1)/3向下取整减去 ( l − 1 ) / 3 (l-1)/3 (l−1)/3向上取整再加1。同理我们可以计算得到除3余0，除3余2的元素的个数。
- 当n=1的时候，答案就是[l,r]中能被3整除的数的个数 b 0 b_0 b0；当n=2的时候，如果第一个位置已经放好，那么第二个位置可以放的数字是有一定的限制的（因为要满足和模3等于0），具体来说，如果第一个位置放的数字除3余0的话，第二位只能放除3余0的那些数；如果第一个位置放的数字除3余1的话，第二位只能放除3余2的那些数，如果第一个位置放的数字除3余2的话，第三位只能放除3余1的那些数。同理可以类推到第 n − 1 n-1 n−1位和第 n n n位的关系，最终的答案就是 a r r [ n ] [ 0 ] arr[n][0] arr[n][0]
- 所以我们可以得到一个递推式
- a r r [ i ] [ 0 ] = a r r [ i − 1 ] [ 0 ] ∗ b 0 + a r r [ i − 1 ] [ 1 ] ∗ b 2 + a r r [ i − 1 ] [ 2 ] ∗ b 1 arr[i][0] = arr[i-1][0] * b_{0} + arr[i-1][1] * b_{2} + arr[i-1][2] * b_{1} arr[i][0]=arr[i−1][0]∗b0+arr[i−1][1]∗b2+arr[i−1][2]∗b1
- a r r [ i ] [ 1 ] = a r r [ i − 1 ] [ 0 ] ∗ b 1 + a r r [ i − 1 ] [ 1 ] ∗ b 0 + a r r [ i − 1 ] [ 2 ] ∗ b 2 arr[i][1] = arr[i-1][0] * b_{1} + arr[i-1][1] * b_{0} + arr[i-1][2] * b_{2} arr[i][1]=arr[i−1][0]∗b1+arr[i−1][1]∗b0+arr[i−1][2]∗b2
- a r r [ i ] [ 2 ] = a r r [ i − 1 ] [ 0 ] ∗ b 2 + a r r [ i − 1 ] [ 1 ] ∗ b 1 + a r r [ i − 1 ] [ 2 ] ∗ b 0 arr[i][2] = arr[i-1][0] * b_{2} + arr[i-1][1] * b_{1} + arr[i-1][2] * b_{0} arr[i][2]=arr[i−1][0]∗b2+arr[i−1][1]∗b1+arr[i−1][2]∗b0
- 然后你就可以用时间复杂度为O(n)，空间复杂度也为O(n)的代码轻松过了。
#include 这是python代码，这是Java代码
你以为结束了么
-
做到这里这个题其实是可以很轻松的通过了，但是如果 n n n特别大，导致你不能再O(n)的时间复杂度内计算出来呢，我们又该如何应对？在找应对方法之前我们先看一个很经典的问题，斐波那契数列数列求第n项。斐波那契数列满足条件 f 0 = 1 f_0=1 f0=1, f 1 = 1 f_1=1 f1=1, f n = f n − 1 + f n − 2 f_n=f_{n-1}+f_{n-2} fn=fn−1+fn−2，当 n > = 2 n>=2 n>=2的时候。你一定知道如何用递归的方式、循环的方式求解斐波那契数列的第n项。但是当n特别大的时候递归&循环的方式是不好使的，会超出内存&时间限制。我们可以尝试把 f n = f n − 1 + f n − 2 f_n=f_{n-1} +f_{n-2} fn=fn−1+fn−2写成矩阵乘法的形式，有
[ 0 1 1 1 ] ∗ [ f n − 2 f n − 1 ] = [ f n − 1 f n ] \left[\begin{matrix}0 & 1 \\ 1 & 1\end{matrix}\right] * \left[\begin{matrix}f_{n-2} \\ f_{n-1}\end{matrix}\right] = \left[\begin{matrix}f_{n-1} \\ f_{n}\end{matrix}\right] [0111]∗[fn−2fn−1]=[fn−1fn] -
进而可以递推得到
[ f n − 1 f n ] = [ 0 1 1 1 ] n − 1 ∗ [ f 0 f 1 ] \left[\begin{matrix}f_{n-1} \\ f_{n}\end{matrix}\right] = \left[\begin{matrix}0 & 1 \\ 1 & 1\end{matrix}\right]^{n-1} * \left[\begin{matrix}f_{0} \\ f_{1}\end{matrix}\right] [fn−1fn]=[0111]n−1∗[f0f1] -
所以我们只要快速求出中间这个矩阵的 n − 1 n-1 n−1次幂就可以了。要求 a b a^b ab的话，如果b为偶数，可以将它分解为两部分 a b = a b / 2 ∗ a b / 2 a^b=a^{b/2} * a^{b/2} ab=ab/2∗ab/2；而如果b为奇数，可以表示为 a b = a b / 2 ∗ a b / 2 ∗ a a^b=a^{b/2} * a^{b/2} * a ab=ab/2∗ab/2∗a。拿b是偶数举例，之前我们计算 a b a^b ab的时候需要O(b)的时间复杂度；而观察到 a b = a b / 2 ∗ a b / 2 a^b=a^{b/2} * a^{b/2} ab=ab/2∗ab/2，当计算得到 a b / 2 a^{b/2} ab/2的时候，只需要再通过一次乘法就可以得到 a b a^b ab。依次类推，就可以以 l o g ( b ) log(b) log(b)的时间复杂度计算得到 a b a^b ab。这种方法就叫做快速幂。
-
当 a a a为矩阵的时候，同样的道理我们可以利用快速幂的思想，以log时间复杂度得到 a b a^b ab。
-
再回到这个题，前面的递推式子同样可以写成矩阵乘法的形式：
[ b 0 b 2 b 1 b 1 b 0 b 2 b 2 b 1 b 0 ] ∗ [ a r r [ n − 1 ] [ 0 ] a r r [ n − 1 ] [ 1 ] a r r [ n − 1 ] [ 2 ] ] = [ a r r [ n ] [ 0 ] a r r [ n ] [ 1 ] a r r [ n ] [ 2 ] ] \left[ \begin{matrix} b_0 & b_2 & b_1 \\ b_1 & b_0 & b_2 \\ b_2 & b_1 & b_0 \end{matrix} \right] * \left[\begin{matrix}arr[n-1][0] \\ arr[n-1][1] \\ arr[n-1][2]\end{matrix}\right] =\left[\begin{matrix}arr[n][0] \\ arr[n][1] \\ arr[n][2]\end{matrix}\right] ⎣⎡b0b1b2b2b0b1b1b2b0⎦⎤∗⎣⎡arr[n−1][0]arr[n−1][1]arr[n−1][2]⎦⎤=⎣⎡arr[n][0]arr[n][1]arr[n][2]⎦⎤ -
依次递推下去的话，可以得到
[ a r r [ n ] [ 0 ] a r r [ n ] [ 1 ] a r r [ n ] [ 2 ] ] = [ b 0 b 2 b 1 b 1 b 0 b 2 b 2 b 1 b 0 ] n − 1 ∗ [ a r r [ 1 ] [ 0 ] a r r [ 1 ] [ 1 ] a r r [ 1 ] [ 2 ] ] \left[\begin{matrix}arr[n][0] \\ arr[n][1] \\ arr[n][2]\end{matrix}\right]= \left[ \begin{matrix} b_0 & b_2 & b_1 \\ b_1 & b_0 & b_2 \\ b_2 & b_1 & b_0 \end{matrix} \right] ^{n-1} * \left[\begin{matrix}arr[1][0] \\ arr[1][1] \\ arr[1][2]\end{matrix}\right] ⎣⎡arr[n][0]arr[n][1]arr[n][2]⎦⎤=⎣⎡b0b1b2b2b0b1b1b2b0⎦⎤n−1∗⎣⎡arr[1][0]arr[1][1]arr[1][2]⎦⎤ -
对于中间这个矩阵的 n − 1 n-1 n−1次幂，可以利用矩阵快速幂计算得到，时间复杂度是O(logn)，这样，即使n很大也能很快解出来了。
#include python代码
import math
n,l,r = map(int, input().split(' '))
zero = math.floor(r/3)-math.ceil(l/3)+1
one = math.floor((r-1)/3)-math.ceil((l-1)/3)+1
two = math.floor((r-2)/3)-math.ceil((l-2)/3)+1
mod=10**9+7
class Matrix:
def __init__(self):
self.m = [[0 for _ in range(3)] for j in range(3)]
def __str__(self):
print(self.m)
return ""
def mul(a, b):
ret = Matrix()
for i in range(3):
for j in range(3):
ret.m[i][j] = 0
for k in range(3):
ret.m[i][j] += a.m[i][k] * b.m[k][j]
ret.m[i][j] %= mod
return ret
def pow(a, b):
ret = Matrix()
ret.m[0][0] = 1
ret.m[1][1] = 1
ret.m[2][2] = 1
while b >= 1:
if b % 2 == 1:
ret = mul(ret, a)
b >>= 1
a = mul(a, a)
return ret
a = Matrix()
a.m[0][0], a.m[1][1], a.m[2][2] = zero, zero, zero
a.m[0][2], a.m[1][0], a.m[2][1] = one, one, one
a.m[0][1], a.m[1][2], a.m[2][0] = two, two, two
ret = pow(a, n-1)
ans = ret.m[0][0] * zero + ret.m[0][1] * one + ret.m[0][2] * two
ans %= mod
print(ans)
总结
- 这个题如果要只是要直接通过的话是很容易的。但是当你拓展开来，发现当面临的问题规模进一步扩大的时候，就得换用别的思路来做题了，这时候就涉及到了更多的知识点，比如矩阵快速幂、矩阵乘法等等。
- 事实上，利用矩阵转移方程求解是一类递推问题的通用解法。面鲸在这里给大家推荐一个博客，可以去看看，相信你会有更多的收获！
- 这里给大家列了几道类似的问题，尝试去切掉它吧。
- https://codeforces.com/contest/1105/problem/C(这个题原题)
- https://leetcode-cn.com/problems/three-steps-problem-lcci/
- https://codeforces.com/problemset/problem/450/B
- https://codeforces.com/contest/1117/problem/D
号外
- 由面鲸公众号组织的每日刷题活动正式开始啦。想要一起起飞的小伙伴欢迎后台联系小编拉你入群呀，与更多的小伙伴一起进步，奥利给！