Artifact Content
Not logged in

Artifact 3402769c07a731ca3bd865c67ae5982d0e2ba42e


#include <iostream>
#include <sstream>
#include <iomanip>
#include <vector>
#include <string>
#include <map>
#include <set>
#include <algorithm>
#include <numeric>
#include <iterator>
#include <complex>
#include <queue>
#include <stack>
#include <cmath>
#include <cassert>
using namespace std;
typedef long long LL;

class ExtendedHappyNumbers
{
public:
	long long calcTheSum(int A, int B, int K) 
	{
		vector<int> happy(9*9*9*9*9*9*7);
//		map<int,int> happy;
		LL sum = 0;
		for(int n=A; n<=B; ++n)
			sum += rec(n, K, happy);
		return sum;
	}

	int S(int K, int n)
	{
		int sum = 0;
		for(int d=1; d<=n; d*=10)
		{
			int v = n/d%10;
			int t = 1;
			for(int i=0; i<K; ++i)
				t *= v;
			sum += t;
		}
		return sum;
	}

	int rec(int n, int K, vector<int>& happy)
	{
		int LOOP = -1;
		stack<int> vs;
		int m = 0x7fffffff;
		for(;;)
		{
			if( happy[n] == LOOP )
				{m = looped(n, K); break;}
			if( happy[n] )
				{m = happy[n]; break;}
			happy[n] = LOOP;
			vs.push(n);
			n = S(K, n);
		}
		while(!vs.empty()) {
			int v = vs.top(); vs.pop();
			happy[v] = m = min(m, v);
		}
		return m;
	}

	int looped(int n, int K)
	{
		int mn = n;
		vector<int> vs(1, n);
		for(int v=S(K,n); v!=n; v=S(K,v)) 
			mn = min(mn, v);
		return mn;
	}
};

// BEGIN CUT HERE
#include <ctime>
double start_time;string timer() { ostringstream os; os << " (" << int((clock()-start_time)/CLOCKS_PER_SEC*1000) << " msec)"; return os.str(); }

template <typename T> string print_array(const vector<T> &V) { ostringstream os; os << "{ "; for (typename vector<T>::const_iterator iter = V.begin(); iter != V.end(); ++iter) os << '\"' << *iter << "\","; os << " }"; return os.str(); }
int verify_case(const long long &Expected, const long long &Received) { if (Expected == Received) cerr << "PASSED" << timer() << endl; else { cerr << "FAILED" << timer() << endl; cerr << "\tExpected: \"" << Expected << '\"' << endl; cerr << "\tReceived: \"" << Received << '\"' << endl; } return 0;}

template<int N> struct Case_ { Case_(){start_time=clock();} };
char Test_(...);
int Test_(Case_<0>) {
	int A = 13; 
	int B = 13; 
	int K = 2; 
	long long RetVal = 1LL; 
	return verify_case(RetVal, ExtendedHappyNumbers().calcTheSum(A, B, K)); }
int Test_(Case_<1>) {
	int A = 1; 
	int B = 5; 
	int K = 2; 
	long long RetVal = 14LL; 
	return verify_case(RetVal, ExtendedHappyNumbers().calcTheSum(A, B, K)); }
int Test_(Case_<2>) {
	int A = 10; 
	int B = 99; 
	int K = 1; 
	long long RetVal = 450LL; 
	return verify_case(RetVal, ExtendedHappyNumbers().calcTheSum(A, B, K)); }
int Test_(Case_<3>) {
	int A = 535; 
	int B = 538; 
	int K = 3; 
	long long RetVal = 820LL; 
	return verify_case(RetVal, ExtendedHappyNumbers().calcTheSum(A, B, K)); }
int Test_(Case_<4>) {
	int A = 100000; 
	int B = 400000; 
	int K = 6; 
	long long RetVal = 5169721292LL; 
	return verify_case(RetVal, ExtendedHappyNumbers().calcTheSum(A, B, K)); }

template<int N> void Run_() { cerr << "Test Case #" << N << "..." << flush; Test_(Case_<N>()); Run_<sizeof(Test_(Case_<N+1>()))==1 ? -1 : N+1>(); }
template<>      void Run_<-1>() {}
int main() { Run_<0>(); }
// END CUT HERE