Sadonkorjuu

In file included from /usr/include/c++/11/pstl/parallel_backend_serial.h:16,
                 from /usr/include/c++/11/pstl/parallel_backend.h:14,
                 from /usr/include/c++/11/pstl/algorithm_impl.h:22,
                 from /usr/include/c++/11/pstl/glue_execution_defs.h:50,
                 from /usr/include/c++/11/execution:32,
                 from input/code.cpp:9:
/usr/include/c++/11/numeric: In instantiation of 'constexpr _Tp std::transform_reduce(_InputIterator, _InputIterator, _Tp, _BinaryOperation, _UnaryOperation) [with _InputIterator = __gnu_cxx::__normal_iterator<std::pair<int, int>*, std::vector<std::pair<int, int> > >; _Tp = long long int; _BinaryOperation = solve(const std::vector<std::vector<std::pair<int, int> > >&, std::vector<short int>, size_t)::<lambda(std::pair<int, int>, size_t)>; _UnaryOperation = __pstl::__internal::__no_op]':
/usr/include/c++/11/pstl/numeric_impl.h:96:33:   required from '_Tp __pstl::__internal::__brick_transform_reduce(_ForwardIte...

#include <vector>
#include <utility>
#include <map>
#include <algorithm>
#include <iostream>
#include <cstdint>
#include <queue>
#include <limits>
#include <execution>

using namespace std;
using LLi = long long int;

static const LLi MAX_DIST = std::numeric_limits<LLi>::max();


auto solve(const vector<vector<pair<int32_t, int32_t>>>& edges, const vector<int16_t> node_types, size_t n) -> size_t
{
    priority_queue<pair<int32_t, int32_t>, std::vector<pair<int32_t, int32_t>>, std::greater<pair<int32_t, int32_t>>> distances;
    vector<bool> is_processed(n + 1, false); 
    vector<pair<int32_t, int32_t>> vertex_attributes(n + 1, make_pair(MAX_DIST, 0));

    for (size_t i = 0; i < node_types.size(); i++) {
        int16_t curr = node_types[i];
        if (curr == 0) {
            vertex_attributes[i] = make_pair(0, i);
            distances.push( vertex_attributes[i]);
        }
    }

    
    auto Relax = [&vertex_attributes]( const pair<int32_t, int32_t>& curr, const int32_t a_neighbour, int32_t weight ) -> void
    {
        if ( curr.first + weight < vertex_attributes[a_neighbour].first )  {
            vertex_attributes[a_neighbour].first = curr.first + weight;
            vertex_attributes[a_neighbour].second = curr.second;
        }
    };


    while ( !(distances.empty()) ) {
        pair<int32_t, int32_t> curr = distances.top();
        distances.pop();

        for ( const pair<int32_t, int32_t>& a_node : edges[curr.second] ) {
            if ( !is_processed[a_node.second] && (node_types[a_node.second] == 1) ) {
                Relax(curr, a_node.second, a_node.first);
                distances.emplace( vertex_attributes[a_node.second].first, a_node.second );
            }
        }
    }

    return reduce(execution::par, vertex_attributes.begin(), vertex_attributes.end(), 0LL, 
    [](pair<int32_t, int32_t> curr, size_t total)
    {
        return total + curr.second;
    });

}



int main()
{
    size_t n = 0;
    
    cin >> n;

    vector<int16_t> node_types(n + 1);
    vector<vector<pair<int32_t, int32_t>>> edges(n + 1);

    for (size_t i = 1; i <= n; i++) {
        edges[i].reserve(n / 10);
    }

    for (size_t i = 1; i <= n; i++) {
        cin >> node_types[i];
    }


    int32_t a = 0;
    int32_t b = 0;
    int32_t dist = 0;
    for (size_t i = 0; i < n - 1; i++) {
        cin >> a >> b >> dist;
        edges[a].push_back(make_pair(dist, b));
        edges[b].push_back(make_pair(dist, a));
        
    }

    cout << solve(edges, node_types, n) << endl;

    return EXIT_SUCCESS;
}