백준 2331번 반복수열 C++

https://www.acmicpc.net/problem/2331

2331번: 반복수열

 

#include <iostream>
#include <string>
#include <vector>
#include <math.h>
#include <algorithm>
#include <utility>
#include <stack>
#include <queue>
#include <math.h>
#include <set>
#include <map>
#include <list>
#include <unordered_map>
#include <unordered_set>
#include <iomanip>
#include <limits.h>

using namespace std;
using int64 = long long;

int64 a;
int p;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    cin >> a >> p;
    vector<int64> v;
    set<int64> s;
    v.push_back(a);
    s.insert(a);
    int ans = 0;
    while (true) {
        int64 currNum = v.back();
        int64 nextNum = 0;
        string str = to_string(currNum);
        for (int i = 0; i < str.length(); i++) {
            nextNum += pow((str[i] - '0'), p);
        }
        auto it = s.find(nextNum);
        if (it == s.end()) {
            s.insert(nextNum);
            v.push_back(nextNum);
        }
        else {
            for (int i = 0; i < v.size(); i++) {
                if (v[i] == nextNum) {
                    ans = i;
                    break;
                }
            }
            break;
        }
    }
    cout << ans;
}