Sadonkorjuu

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

using namespace std;

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

    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;
        vector<int> visited;

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



int main()
{
    // 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);
    }

    // ALG
    queue<SearchState> q = queue<SearchState>();
    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, visited));
        }
    }

    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(state.visited.begin(), state.visited.end(), connection.destination->cityID)) { continue; }

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

            state.visited.push_back(connection.destination->cityID);
            SearchState newState = SearchState(connection.destination, newCost, vector<int>(state.visited));
            q.push(newState);
        }
    }

    long totalLength = 0;
    for (int i = 0; i < cityCount; i++){
        if(cities[i].isField){
            totalLength += cities[i].shortestPath;
        }
    }

    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: TIME LIMIT EXCEEDED

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