TopCoder SRM 334 Div1 Medium ExtendedHappyNumbers

問題

Nに対してSk(N)を、
Σ(Nの各桁)^kと定義する。
Nに対して、Nの幸福度を、数列N,Sk(N),Sk(Sk(N)),...,の最小値とする。

与えられた整数A,Bに対して、
A≦i≦Bなるすべてのに対するiの幸福度の和を求めよ。

制約条件

A,B≦1000000
k≦6

方針

メモ化してiに対する幸福度を求める。
Sk(i)がループしたら、ループ部分の幸福度はループ部分の最小値となる。

ループの検出は、フラグをたてながら深さ優先探索をすればよい。
いままで訪問した数字をスタックに入れておいて、ループの部分の数字を書き換えるのに使う。

ソースコード

const int MX=4000000;
int k,dp[MX];
int sz,st[MX];
bool v[MX];
int rec(int c){
  if(v[c]){
    int tmp=dp[c];
    for(int i=sz-1;i>=0&&st[i]!=c;i--)tmp=min(tmp,st[i]);
    for(int i=sz-1;i>=0&&st[i]!=c;i--)dp[st[sz-1]]=tmp;
    return dp[c]=tmp;
  }
  if(dp[c]>=0)return dp[c];
  v[c]=1;
  dp[c]=c; st[sz++]=c;
  int nxt=0;
  for(int m=c;m;m/=10){
    int t=m%10, s=1;
    rep(i,k)s*=t;
    nxt+=s;
  }
  return dp[c]=min(dp[c],rec(nxt));
}

class ExtendedHappyNumbers {
  public:
  long long calcTheSum(int A, int B, int K) {
    memset(dp,-1,sizeof(dp));
    memset(v,0,sizeof(v));
    k=K;
    ll ans=0;
    for(int i=A;i<=B;i++){
      sz=0;
      ans+=rec(i);
      rep(j,sz)v[st[j]]=0;
    }
    return ans;
  }
};