star2dust

力扣链接：https://leetcode-cn.com/leetbook/read/didiglobal2/e7hh2i/

给定三种类型的小球 P、Q、R，每种小球的数量分别为 np、nq、nr 个。现在想将这些小球排成一条直线，但是不允许相同类型的小球相邻，问有多少种排列方法。如果无法组合出合适的结果，则输出 0。

格式：

输入：一行以空格相隔的三个数，分别表示为 np，nq，nr。
输出：排列方法的数量。

示例：

输入：2 1 1
输出：6
解释：如若 np=2，nq=1，nr=1 则共有 6 种排列方式：PQRP，QPRP，PRQP，RPQP，PRPQ 以及 PQPR。

解法：通过dfs递归求解。

#include<bits/stdc++.h>
using namespace std;

int dfs(int last_color, int np, int nq, int nr){
    // if balls of any color run out
    if (np<0||nq<0||nr<0){
        return 0;
    }
    // if only one ball remains 
    if (np+nq+nr==1){
        if (last_color==0) return np;
        if (last_color==1) return nq;
        if (last_color==2) return nr;
    }else{ 
        // back to n-1
        if (last_color==0) return dfs(1,np-1,nq,nr)+dfs(2,np-1,nq,nr);
        if (last_color==1) return dfs(0,np,nq-1,nr)+dfs(2,np,nq-1,nr);
        if (last_color==2) return dfs(0,np,nq,nr-1)+dfs(1,np,nq,nr-1);
    }
    return 0;
}

int main(){
    int np, nq, nr;
    cin>>np>>nq>>nr;
    // 3 colors 
    int ans = dfs(0,np,nq,nr)+dfs(1,np,nq,nr)+dfs(2,np,nq,nr);
    cout<<ans<<endl;
    return 0;
}

很不幸，上面的解法超时了，测试用例通过3/5，那么加个dp数组进行记忆化搜索。

#include<bits/stdc++.h>
using namespace std;
// 用int不够，改为long long
typedef long long LL;

LL p,q,r;

inline LL I(int c,int x,int y,int z){
    return c*(p+1)*(q+1)*(r+1)+x*(q+1)*(r+1)+y*(r+1)+z;
}

vector<LL> dp;

LL find(int c,int i,int j,int k){
    if(i<0||j<0||k<0) return 0;
    LL idx=I(c,i,j,k);
    if(dp[idx]==-1){
        dp[idx]=0;
        if(c==0) dp[idx]+=find(1,i-1,j,k)+find(2,i-1,j,k);
        if(c==1) dp[idx]+=find(0,i,j-1,k)+find(2,i,j-1,k);
        if(c==2) dp[idx]+=find(0,i,j,k-1)+find(1,i,j,k-1);
    }
    return dp[idx];
}

int main(){
    cin>>p>>q>>r;
    dp.resize(3*(p+1)*(q+1)*(r+1));
    fill(dp.begin(),dp.end(),-1);
    dp[I(0,1,0,0)]=1;
    dp[I(1,0,1,0)]=1;
    dp[I(2,0,0,1)]=1;
    cout<<find(0,p,q,r)+find(1,p,q,r)+find(2,p,q,r)<<endl;
    return 0;
}