ARC117 C - Tricolor Pyramid

题意

给出一个长为 $n$ 的数组，每个元素为 W，B 或 R，表示每个块的颜色。通过这个块，可以输出一个三角形。生成规则：如果相邻两个块的颜色相同，则得到一个同色的块；如果两个块的颜色不同，则得到没有出现过的颜色。如下图所示：

求最上面的块的颜色。
数据范围：$2\le N\le 4\times 10^5$

分析

主要在如何分析转化颜色之间的转换关系，可以令三种颜色分别为：0，1，2，通过加法生成新的块的颜色。具体见下图：

其中，$d_i$ 表示该方块颜色代表的数字。思路可以参考视频。
最终的结果为：
$$(-1{)}^{N-1}\sum _{i=1}^{N-1}{d}_i\times C(N-1,i-1)\mathrm{m}\mathrm{o}\mathrm{d}3$$
问题的难点转化为如何在 $n,m$比 $p$ 大的情况下计算 $C(n,m)\mathrm{m}\mathrm{o}\mathrm{d}3$，由此，可以联想到 $\text{Lucas}$ 定理，最后解决问题。

代码

#include <cstdio>
#include <algorithm>
#include <cmath>
#include <iostream>
using namespace std;
const int N=4e5+5;
char s[N];
int n;
int solve(int a,int b,int p){
    if(b>a) return 0;
    b=min(b,a-b);
    int x=1,y=1;
    for(int i=1;i<=b;i++){
        y=y*i%p;
        x=x*(a-i+1)%p;
    }
    return x*y%p;
}
int Lucas(int a,int b,int p){
    return b?solve(a%p,b%p,p)*Lucas(a/p,b/p,p)%p:1;
}
int main(){
    int ans=0;
    scanf("%d",&n);
    scanf("%s",s+1);
    for(int i=1;i<=n;i++){
        int d=0;
        if(s[i]=='W') d=0;
        else if(s[i]=='R') d=1;
        else if(s[i]=='B') d=2;
        ans=(ans+d*Lucas(n-1,i-1,3)%3)%3;//cout<<i<<":"<<ans<<endl;
    }
    if((n-1)&1) ans=(-ans%3+3)%3;
    char ch;
    if(ans==0) ch='W';
    else if(ans==1) ch='R';
    else if(ans==2) ch='B';
    printf("%c\n",ch);
    return 0;
}