Untitled

#include <iostream>
#include <vector>
#include <queue>
using namespace std;
#define MAX_N 5000
#define INF 1000000

struct Road{
	int begin;
	int end;
	int distance;
};

struct City{
	int id; //khach san
	int brand;
	vector<Road> roads;
};

struct Node {
	int id;
	int cost;
};

struct Comparator{
	inline bool operator() (const Node& n1, const Node& n2){
		return n1.cost > n2.cost; //Node co cost nho hon thi duoc uu tien
	}
};

City cities[MAX_N];
int dist[MAX_N];
int n;

void init(int N, int mBrands[])
{
	for(int i = 0; i < N; i++){
		cities[i].id = i;
		cities[i].brand = mBrands[i];
		cities[i].roads.clear();
	}
	n = N;
}

void connect(int mHotelA, int mHotelB, int mDistance)
{
	Road roadA;
	roadA.begin = mHotelA;
	roadA.end = mHotelB;
	roadA.distance = mDistance;
	cities[mHotelA].roads.push_back(roadA);

	/*City* cityA = &cities[mHotelA];
	cityA->roads.push_back(roadA);*/

	Road roadB;
	roadB.begin = mHotelB;
	roadB.end = mHotelA;
	roadB.distance = mDistance;
	cities[mHotelB].roads.push_back(roadB);
}

int merge(int mHotelA, int mHotelB)
{
	int brandA = cities[mHotelA].brand;
	int brandB = cities[mHotelB].brand;
	int cnt = 0;
	for(int i = 0; i < n; i++){
		if(cities[i].brand == brandB){
			cities[i].brand = brandA;
		}
		if(cities[i].brand == brandA){
			cnt++;
		}
	}
    return cnt;
}

int move(int mStart, int mBrandA, int mBrandB)
{
	for(int i = 0; i < n; i++){
		dist[i] = INF;
	}

	priority_queue<Node, vector<Node>, Comparator> pq;
	Node first;
	first.cost = 0;
	first.id = mStart;
	pq.push(first);
	dist[mStart] = 0;
	int best1 = INF;
	int best2 = INF;
	while(!pq.empty()){
		Node current = pq.top();
		pq.pop();
		if(current.id != mStart){
			if(mBrandA == mBrandB){
				if(cities[current.id].brand == mBrandA){
					if(best1 == INF) best1 = current.cost;

					else if(best2 == INF) {
						best2 = current.cost;
						return best1 + best2;
					}
				}
				
			}
			else{
				if(cities[current.id].brand == mBrandA){
					if(best1 == INF) best1 = current.cost;
				}
				else if(cities[current.id].brand == mBrandB) {
					if(best2 == INF) best2 = current.cost;
				}
				if(best1 != INF && best2 != INF) 
					return best1 + best2;
			}

		}
		for(int i =0 ; i < cities[current.id].roads.size(); i++){
			Road road = cities[current.id].roads[i];
			if(current.cost + road.distance < dist[road.end]){
				dist[road.end] = current.cost + road.distance ;
				Node newNode;
				newNode.id = road.end;
				newNode.cost = dist[road.end];
				pq.push(newNode);
			}
		}
	}

	return 0;
}