Building Teams

#include <iostream>
#include <algorithm>
#include <queue>

using namespace std;

int main() {
    int n, m;
    cin >> n >> m;

    vector<int> dist(n+1,0);

    vector<pair<int,int>> friendships;

    for(int i = 0; i < m; i++) {
        int a, b;
        cin >> a >> b;
        friendships.push_back(make_pair(a,b));
    }

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

    for(size_t i = 0; i < friendships.size(); i++) {
       if(dist[friendships[i].first] == 1 && dist[friendships[i].second] == 0) {
           dist[friendships[i].second] = 2;
       }
       else if(dist[friendships[i].first] == 2 && dist[friendships[i].second] == 0) {
           dist[friendships[i].second] = 1;
       }
       else if(dist[friendships[i].first] == 0 && dist[friendships[i].second] == 1) {
           dist[friendships[i].first] = 2;
       }
       else if(dist[friendships[i].first] == 0 && dist[friendships[i].second] == 2) {
           dist[friendships[i].first] = 1;
       }
       else if(dist[friendships[i].first] == 0 && dist[friendships[i].second] == 0) {
           dist[friendships[i].first] = 1;
           dist[friendships[i].second] = 2;
           //cout << friendships[i].first << " " << friendships[i].second << endl;

       }
       else if(dist[friendships[i].first] == 2 && dist[friendships[i].second] == 1) {
           cout << "IMPOSSIBLE" << endl;
           return 0;
       }
       else if(dist[friendships[i].first] == 1 && dist[friendships[i].second] == 2) {
           cout << "IMPOSSIBLE" << endl;
           return 0;
       }
       
    }

    for(int i = 1; i <= n; i++) {
        if(dist[i] == 0 || dist[i] == 1) cout << 1 << " ";
        else cout << 2 << " ";
        //cout << dist[i] << endl;
    }

}

Test 1

Verdict: WRONG ANSWER

Test 2

Verdict: WRONG ANSWER

Test 3

Verdict: WRONG ANSWER

Test 4

Verdict: WRONG ANSWER

Test 5

Verdict: ACCEPTED

Test 6

Verdict: WRONG ANSWER

Test 7

Verdict: WRONG ANSWER

Test 8

Verdict: WRONG ANSWER

Test 9

Verdict: WRONG ANSWER

Test 10

Verdict: ACCEPTED

Test 11

Verdict: WRONG ANSWER

Test 12

Verdict: ACCEPTED