Untitled

package HotelTravel;

import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Queue;

class UserSolution {
	Hotel[] hotels;

	void init(int N, int mBrands[]) {
		hotels = new Hotel[N];
		for (int i = 0; i < N; i++) {
			hotels[i] = new Hotel();
			hotels[i].brand = mBrands[i];
		}
	}

	void connect(int mHotelA, int mHotelB, int mDistance) {
		hotels[mHotelA].roads.put(hotels[mHotelB], mDistance);
		hotels[mHotelB].roads.put(hotels[mHotelA], mDistance);
	}

	int merge(int mHotelA, int mHotelB) {
		int a = hotels[mHotelA].brand;
		int b = hotels[mHotelB].brand;
		int count = 0;
		for (Hotel hotel : hotels) {
			if (hotel.brand == b || hotel.brand == a) {
				hotel.brand = a;
				++count;
			}
		}
		return count;
	}

	int move(int mStart, int mBrandA, int mBrandB) {
		for (Hotel hotel : hotels) { // ?
			hotel.distance = Integer.MAX_VALUE;
		}
		Queue<Hotel> queue = new PriorityQueue<>((a, b) -> a.distance
				- b.distance);
		hotels[mStart].distance = 0;
		queue.add(hotels[mStart]);

		int result = 0;
		int count = 0;
		while (true) {
			Hotel hotel = queue.poll(); // lấy hotel có khoảng cách ngắn nhất
										// tới hotel xuất phát
			if (hotel.distance > 0
					&& (hotel.brand == mBrandA || hotel.brand == mBrandB)) {
				if (hotel.brand == mBrandA) {
					mBrandA = -1;
				} else {
					mBrandB = -1;
				}
				result += hotel.distance;
				if (++count == 2) {
					return result;
				}
			}

			for (Hotel next : hotel.roads.keySet()) {
				int road = hotel.roads.get(next);
				if (hotel.distance + road < next.distance) {
					queue.remove(next);
					next.distance = hotel.distance + road;
					queue.add(next);
				}
			}
		}
	}

	class Hotel {
		int brand;
		int distance; // khoảng cách min từ vị trí start tới hotel này
		Map<Hotel, Integer> roads = new HashMap<>();
		// mỗi hotel có một map, lưu khoảng cách đường tới các hotel bên cạnh
		// dùng để connect, thay cho arraylist
	}

}