Sadonkorjuu

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
#include <chrono>

using namespace std;

class Connection;
class City {
public:
    int cityID;
    bool isField;
    bool hasFieldConnection = false;
    vector<Connection> connections;
    int shortestPath = 999999999;
    vector<int> visitedBy;
    
    City(){
        this->cityID = -1;
        this->isField = false;
        connections = vector<Connection>();
        visitedBy = vector<int>();
    }
    City(int cityID, bool isField){
        this->cityID = cityID;
        this->isField = isField;
        connections = vector<Connection>();
        visitedBy = vector<int>();
    }

    void CreateConnection(City& destination, int cost);

};

class Connection
{
    public:
        City* destination;
        int cost;

        Connection(City* destination, int cost)
        {
            this->destination = destination;
            this->cost = cost;
        }
};

void City::CreateConnection(City &destination, int cost){
    connections.push_back(Connection(&destination, cost));
    if(destination.isField){
        this->hasFieldConnection = true;
    }
}

class SearchState{
    public:
        City* city;
        int totalCost;
        int startingPort;

        SearchState(){}
        SearchState(City* city, int totalCost, int startingPort){
            this->city = city;
            this->totalCost = totalCost;
            this->startingPort = startingPort;
        }
};



int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    // INPUT
    int cityCount;
    cin >> cityCount;
    vector<City> cities(cityCount);
    for (int i = 0; i < cityCount; ++i) {
        int type;
        cin >> type;
        cities[i] = City(i, type == 1);
    }
    for(int i = 0; i < cityCount - 1; i++){
        int a, b, c;
        cin >> a >> b >> c;
        cities[a - 1].CreateConnection(cities[b - 1], c);
        cities[b - 1].CreateConnection(cities[a - 1], c);
    }

    /*int cityCount = 201111;
    vector<City> cities(cityCount);
    int fieldCount = 0;
    int k = 1.2;
    for(int i = 0; i < cityCount; i++){
        cities[i] = City(i, i < (cityCount / k));
        if (i < (cityCount / k)){
            fieldCount++;
        }
    }
    for(int i = 0; i < cityCount - 1; i++){
        int randomIndex;
        do {
            randomIndex = rand() % cityCount;
        }
        while (randomIndex == i);
        cities[i].CreateConnection(cities[randomIndex], 1);
        cities[randomIndex].CreateConnection(cities[i], 1);
    }

    cout << "FIELDS: " << fieldCount << "\n";*/

    //auto start = std::chrono::high_resolution_clock::now();

    // ALG
    queue<SearchState> q = queue<SearchState>();
    long fieldCount = 0;
    for (int i = 0; i < cityCount; i++){
        if(!cities[i].isField && cities[i].hasFieldConnection){
            vector<int> visited = vector<int>(1);
            visited[0] = i;
            q.push(SearchState(&cities[i], 0, i));
        }
        else if (cities[i].isField){
            fieldCount += 1;
        }
    }
    
    long totalLength = fieldCount * 999999999;
    while (q.size() > 0)
    {
        SearchState state = q.front(); q.pop();
        if (state.city->shortestPath < state.totalCost){
            continue;
        }
        for (Connection connection : state.city->connections){
            if (!connection.destination->isField) { continue; }
            if (count(connection.destination->visitedBy.begin(), connection.destination->visitedBy.end(), state.startingPort)) { continue; }

            int newCost = state.totalCost + connection.cost;
            if(newCost < connection.destination->shortestPath){
                totalLength -= connection.destination->shortestPath - newCost;
                connection.destination->shortestPath = newCost;
            } else { continue; }

            connection.destination->visitedBy.push_back(state.startingPort);
            SearchState newState = SearchState(connection.destination, newCost, state.startingPort);
            q.push(newState);
        }
    }

    //auto end = std::chrono::high_resolution_clock::now();
    //cout << "TIME: " << std::chrono::duration_cast<std::chrono::milliseconds>(end-start).count() / 1000.0 << "\n"; 
    
    cout << totalLength;

}

Test 1

Group: 1, 2

Verdict: ACCEPTED

Test 2

Group: 1, 2

Verdict: ACCEPTED

Test 3

Group: 1, 2

Verdict: ACCEPTED

Test 4

Group: 1, 2

Verdict: ACCEPTED

Test 5

Group: 1, 2

Verdict: ACCEPTED

Test 6

Group: 1, 2

Verdict: ACCEPTED

Test 7

Group: 2

Verdict: ACCEPTED

Test 8

Group: 1, 2

Verdict: ACCEPTED

Test 9

Group: 2

Verdict: ACCEPTED

Test 10

Group: 1, 2

Verdict: ACCEPTED

Test 11

Group: 2

Verdict: ACCEPTED

Test 12

Group: 2

Verdict: ACCEPTED

Test 13

Group: 2

Verdict: ACCEPTED

Test 14

Group: 2

Verdict: ACCEPTED

Test 15

Group: 1, 2

Verdict: ACCEPTED

Test 16

Group: 1, 2

Verdict: ACCEPTED

Test 17

Group: 1, 2

Verdict: ACCEPTED

Test 18

Group: 1, 2

Verdict: ACCEPTED

Test 19

Group: 1, 2

Verdict: ACCEPTED

Test 20

Group: 1, 2

Verdict: ACCEPTED

Test 21

Group: 2

Verdict: ACCEPTED

Test 22

Group: 2

Verdict: ACCEPTED

Test 23

Group: 2

Verdict: ACCEPTED

Test 24

Group: 1, 2

Verdict: ACCEPTED

Test 25

Group: 2

Verdict: ACCEPTED

Test 26

Group: 1, 2

Verdict: ACCEPTED

Test 27

Group: 2

Verdict: ACCEPTED

Test 28

Group: 1, 2

Verdict: ACCEPTED

Test 29

Group: 2

Verdict: ACCEPTED

Test 30

Group: 1, 2

Verdict: ACCEPTED

Test 31

Group: 2

Verdict: ACCEPTED