가티있는블로그

[SWEA] 5521.상원이의 생일파티

2020. 3. 25. 14:54 | 프로그래밍/코딩테스트 문제

SW Expert Academy에서 5521. 상원이의 생일파티 문제를 해결했다.

중복 값을 제외할 수 있는 set을 사용하여 문제를 쉽게 해결할 수 있었다.

  • s에는 상원이의 친한친구를 저장한다.
  • s1에는 상원이의 친한친구의 친한친구를 저장한다.
  • s에 s1의 원소들을 합쳐서 s의 크기를 반환하여서 정답을 구하였다.
#include<iostream>
#include<algorithm> 
#include<set>
#include<vector>

using namespace std;

int main(int argc, char** argv) {
    int test_case;
    int T;
    cin >> T;

    for (test_case = 1; test_case <= T; ++test_case) {
        int n, m;
        cin >> n >> m;
        //상원이의 친한친구 저장
        set<int> s;
        //친한친구의 친한친구
        set<int> s1;
        //친한친구목록
        vector<pair<int, int>> v;
        for (int i = 0; i < m; i++) {
            int n1;
            int n2;
            cin >> n1>>n2;
            if (n1 == 1) {
                s.insert(n2);
            }
            else v.push_back((make_pair(n1, n2)));
        }

        sort(v.begin(), v.end());

        set<int>::iterator iter;
        for (int i = 0; i < v.size(); i++) {       
            iter = s.find(v[i].first);
            if (iter != s.end()) {
                s1.insert(v[i].second);
            }
            else {
                iter = s.find(v[i].second);
                if (iter != s.end()) {
                    s1.insert(v[i].first);
                }
            }
        }

        for (iter = s1.begin(); iter != s1.end(); ++iter) {
            s.insert(*iter);
        }
        cout << "#" << test_case << " " << s.size() << "\n";
    }
    return 0;//정상종료시 반드시 0을 리턴해야합니다.
}