笔试 | 字节跳动2021秋招第三场 数数组题目详解&扩展

文章首发于公众号面鲸，欢迎关注。另外由面鲸组织的每天一道高频面试题刷题挑战活动正在进行中，快来挑战吧～

数数组

题目描述

• 现在要生成一个数列，满足以下条件
  • 数列的长度为n，且每个数字的大小$a[i]$满足$l<=a[i]<=r$
  • 数组的和要能被3整除
• 现在给定三个数$n,l,r$，可否求出满足以上条件的数组个数，因为数组数量比较大，最终结果请mod 1e9+7

分析

• 遇到这种问题，一看是mod 1e9+7，那么结果一定很大，应该是有一定的递推式子的。
• 令$b_0$表示[l,r]中除3余0的数的个数，$b_1$表示[l,r]中除3余1的数的个数，$b_2$表示[l,r]中除3余2的个数；令$arr[i][x]$表示前i位除3余x的方案数，有$arr[1][x]=b_x$。一种最直接计算$b_x$的方法是遍历[l,r]中的所有整数，挨个统计。但是注意到[l,r]的范围可能很大，因此不能直接这么做。那有什么办法可以快速计算得到这些数呢。如果我们要统计除3余1的数的个数，那么这个数可以表示为$3k+1$，进而有$l<=3k+1<=r$，也就是$(l-1)/3<=k<=(r-1)/3$，又因为k必须是整数，所以k的个数为$(r-1)/3$向下取整减去$(l-1)/3$向上取整再加1。同理我们可以计算得到除3余0，除3余2的元素的个数。
• 当n=1的时候，答案就是[l,r]中能被3整除的数的个数$b_0$；当n=2的时候，如果第一个位置已经放好，那么第二个位置可以放的数字是有一定的限制的（因为要满足和模3等于0），具体来说，如果第一个位置放的数字除3余0的话，第二位只能放除3余0的那些数；如果第一个位置放的数字除3余1的话，第二位只能放除3余2的那些数，如果第一个位置放的数字除3余2的话，第三位只能放除3余1的那些数。同理可以类推到第$n-1$位和第$n$位的关系，最终的答案就是$arr[n][0]$
• 所以我们可以得到一个递推式
  • $arr[i][0]=arr[i-1][0]\ast b_0+arr[i-1][1]\ast b_2+arr[i-1][2]\ast b_1$
  • $arr[i][1]=arr[i-1][0]\ast b_1+arr[i-1][1]\ast b_0+arr[i-1][2]\ast b_2$
  • $arr[i][2]=arr[i-1][0]\ast b_2+arr[i-1][1]\ast b_1+arr[i-1][2]\ast b_0$
• 然后你就可以用时间复杂度为O(n)，空间复杂度也为O(n)的代码轻松过了。

#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
const ll mod=1e9+7;
 
int main() {
    ll n,l,r; cin>>n>>l>>r;
    ll a1,a2,a0;
    a1=floor((long double)(r-1)/3.00)-ceil((long double)(l-1)/3.00)+1;
    a2=floor((long double)(r-2)/3.00)-ceil((long double)(l-2)/3.00)+1;
    a0=floor((long double)(r-3)/3.00)-ceil((long double)(l-3)/3.00)+1;
 
    vector<vector<ll>> dp(n,vector<ll>(3,0));
 
    dp[0][0]=a0; dp[0][1]=a1; dp[0][2]=a2;
 
    for(ll i=1;i<n;i++) {
        dp[i][0]=((dp[i-1][1]*a2)%mod+(dp[i-1][0]*a0)%mod+(dp[i-1][2]*a1)%mod)%mod;
        dp[i][1]=((dp[i-1][0]*a1)%mod+(dp[i-1][1]*a0)%mod+(dp[i-1][2]*a2)%mod)%mod;
        dp[i][2]=((dp[i-1][0]*a2)%mod+(dp[i-1][1]*a1)%mod+(dp[i-1][2]*a0)%mod)%mod;
    }
    cout<<dp[n-1][0]%mod;
 
    return 0;
}

这是python代码，这是Java代码

---

你以为结束了么

• 做到这里这个题其实是可以很轻松的通过了，但是如果$n$特别大，导致你不能再O(n)的时间复杂度内计算出来呢，我们又该如何应对？在找应对方法之前我们先看一个很经典的问题，斐波那契数列数列求第n项。斐波那契数列满足条件$f_0=1$, $f_1=1$, $f_n=f_{n-1}+f_{n-2}$，当$n>=2$的时候。你一定知道如何用递归的方式、循环的方式求解斐波那契数列的第n项。但是当n特别大的时候递归&循环的方式是不好使的，会超出内存&时间限制。我们可以尝试把$f_n=f_{n-1}+f_{n-2}$写成矩阵乘法的形式，有
  $$\left[\begin{array}{cc}0& 1\\ 1& 1\end{array}\right]\ast \left[\begin{array}{c}{f}_{n-2}\\ f_{n-1}\end{array}\right]=\left[\begin{array}{c}{f}_{n-1}\\ f_n\end{array}\right]$$

• 进而可以递推得到
  $$\left[\begin{array}{c}{f}_{n-1}\\ f_n\end{array}\right]={\left[\begin{array}{cc}0& 1\\ 1& 1\end{array}\right]}^{n-1}\ast \left[\begin{array}{c}{f}_0\\ f_1\end{array}\right]$$

• 所以我们只要快速求出中间这个矩阵的$n-1$次幂就可以了。要求$a^b$的话，如果b为偶数，可以将它分解为两部分$a^b=a^{b/2}\ast a^{b/2}$；而如果b为奇数，可以表示为$a^b=a^{b/2}\ast a^{b/2}\ast a$。拿b是偶数举例，之前我们计算$a^b$的时候需要O(b)的时间复杂度；而观察到$a^b=a^{b/2}\ast a^{b/2}$，当计算得到$a^{b/2}$的时候，只需要再通过一次乘法就可以得到$a^b$。依次类推，就可以以$log(b)$的时间复杂度计算得到$a^b$。这种方法就叫做快速幂。

• 当$a$为矩阵的时候，同样的道理我们可以利用快速幂的思想，以log时间复杂度得到$a^b$。

• 再回到这个题，前面的递推式子同样可以写成矩阵乘法的形式：
  $$\left[\begin{array}{ccc}{b}_0& b_2& b_1\\ b_1& b_0& b_2\\ b_2& b_1& b_0\end{array}\right]\ast \left[\begin{array}{c}arr[n-1][0]\\ arr[n-1][1]\\ arr[n-1][2]\end{array}\right]=\left[\begin{array}{c}arr[n][0]\\ arr[n][1]\\ arr[n][2]\end{array}\right]$$

• 依次递推下去的话，可以得到
  $$\left[\begin{array}{c}arr[n][0]\\ arr[n][1]\\ arr[n][2]\end{array}\right]={\left[\begin{array}{ccc}{b}_0& b_2& b_1\\ b_1& b_0& b_2\\ b_2& b_1& b_0\end{array}\right]}^{n-1}\ast \left[\begin{array}{c}arr[1][0]\\ arr[1][1]\\ arr[1][2]\end{array}\right]$$

• 对于中间这个矩阵的$n-1$次幂，可以利用矩阵快速幂计算得到，时间复杂度是O(logn)，这样，即使n很大也能很快解出来了。

#include <bits/stdc++.h>

using namespace std;

#define ll long long
const ll mod = 1e9 + 7;
struct Matrix {
    ll m[3][3];
    Matrix() {
        memset(m, 0, sizeof(m));
    }
};

Matrix mul(Matrix a, Matrix b) {
    Matrix c;
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < 3; ++j) {
            c.m[i][j] = 0;
            for (int k = 0; k < 3; ++k) {
                c.m[i][j] += a.m[i][k] * b.m[k][j];
                c.m[i][j] %= mod;
            }
        }
    }
    return c;
}

Matrix pow(Matrix a, ll b) {
    Matrix ret; ret.m[0][0] = ret.m[1][1] = ret.m[2][2] = 1;
    while (b) {
        if (b&1) {
            ret = mul(ret, a);
        }
        b >>= 1;
        a= mul(a, a);
    }
    return ret;

}

int main() {
    ll n, l, r; cin >> n >> l >> r;
    ll a1=floor((long double)(r-1)/3.00)-ceil((long double)(l-1)/3.00)+1;
    ll a2=floor((long double)(r-2)/3.00)-ceil((long double)(l-2)/3.00)+1;
    ll a0=floor((long double)(r-3)/3.00)-ceil((long double)(l-3)/3.00)+1;

    Matrix a;
    a.m[0][0] = a.m[1][1] = a.m[2][2] = a0;
    a.m[0][2] = a.m[1][0] = a.m[2][1] = a1;
    a.m[0][1] = a.m[1][2] = a.m[2][0] = a2;

    Matrix ret = pow(a, n-1);
    ll ans = ret.m[0][0] * a0 % mod + ret.m[0][1] * a1 % mod + ret.m[0][2] * a2 % mod;
    ans %= mod;
    cout << ans << endl;
    return 0;
}

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

号外

• 由面鲸公众号组织的每日刷题活动正式开始啦。想要一起起飞的小伙伴欢迎后台联系小编拉你入群呀，与更多的小伙伴一起进步，奥利给！