Problema 2 OJI 11-12 2017

#include <fstream>
#include <set>
#include <algorithm>
#include <map>
#include <queue>
#include <vector>
#include <cstdlib>
#include <string>
#include <stack>
#include <utility>
#define NMAX 32000 // DE SCHIMBAT
#define INF 0x3f3f3f3f
#define HASHMOD 333011
#define pb push_back

using namespace std;

ifstream fin("ninjago.in");
ofstream fout("ninjago.out");

int p[NMAX],n,m,dist[NMAX],apm,coridoare,euri;
vector<pair<int, pair<int, int> > > G[NMAX];
pair<pair<int,int>, pair<int,int> > muchii[NMAX];
char s[5];
bool viz[NMAX];

int BFS() {
    int i,x,nr=0,y;

    queue<int> Q;
    viz[1]=1;
    Q.push(1);
    while(!Q.empty()) {
        x=Q.front();
        Q.pop();
        ++nr;

        for(i=0;i<G[x].size();++i) {
            y=G[x][i].first;
            if(viz[y]==0 && G[x][i].second.second==0) {
                viz[y]=1;
                Q.push(y);
            }
        }
    }

    return nr;
}

int findSet(int x) {
    if(p[x]==x) return x;

    return p[x]=findSet(p[x]);
}

bool inSameSet(int x, int y) {
    return findSet(x)==findSet(y);
}

void unionSet(int x, int y) {
    int i=findSet(x);
    int j=findSet(y);

    p[i]=j;
}

void kruskal() {
    int i,x,y,cost,c,e;

    for(i=1;i<=n;++i) p[i]=i;
    sort(muchii+1,muchii+m+1);

    for(i=1;i<=m;++i) {
        c=muchii[i].first.first;
        e=muchii[i].first.second;
        x=muchii[i].second.first;
        y=muchii[i].second.second;

        if(!inSameSet(x,y)) {
            unionSet(x,y);
            coridoare+=(c!=0);
            euri+=c;
            apm+=e;
        }
    }
}

int main() { // NU LASA AFISARI PE ECRAN
    int test,i,j,x,y,p5,cost,hasE;

    fin>>test;
    fin>>n>>m;

    for(i=1;i<=m;++i) {
        fin>>x>>y>>s;

        hasE=cost=0;
        p5=1;
        for(j=0;j<4;++j) {
            if(s[j]=='E') ++hasE; //calculez numarul de E-uri pe fiecare muchie
            else cost+=(s[j]-'A'+1)*p5; // puterea necesara pentru a distruge restul caracterelor
            p5*=5;
        }

        G[x].pb({y,{cost,hasE}});
        G[y].pb({x,{cost,hasE}});
        muchii[i]={{hasE,cost},{x,y}};
    }

    if(test==1) fout<<BFS();
    else {
        kruskal();
        if(test==2) fout<<coridoare<<'\n'<<euri;
        else fout<<apm;
    }

    fin.close();
    fout.close();

    return 0;
}