Tietoverkko

#include <iostream>
#include <string>
#include <math.h>
#include <vector>
#include <bits/stdc++.h>
#include <algorithm>
using namespace std;
 
struct Connection {
    long destination, speed, origin;
    long originCount = -1, destCount = -1;
    bool disabled = false;
 
    Connection(long speed, long dest, long origin) {
        this->destination = dest;
        this->speed = speed;
        this->origin = origin;
    }
};
 
long calcSize(long origin, long current, vector<Connection*> nodes[]) {
    long sum = 1;
 
    for (long i = 0; (unsigned) i < nodes[current].size(); i++) {
        Connection* con = nodes[current].at(i);
        if (!con->disabled) {
            long dest = con->destination;
 
            if (current == dest) {
                dest = con->origin;
            }
 
            if (origin == dest) {
                continue;
            }
 
            long count = calcSize(current, dest, nodes);
            sum += count;
        }
    }
 
    return sum;
}
 
bool compareConnections(Connection* c1, Connection* c2) {
    return c1->speed < c2->speed;
}
 
long solve(vector<Connection*> nodes[], vector<Connection*> connections) {
    
    sort(connections.begin(), connections.end(), compareConnections);
 
    long res = 0;
    for (long i = 0; (unsigned) i < connections.size(); i++) {
        Connection* con = connections.at(i);
        long a = calcSize(con->origin, con->destination, nodes);
        long b = calcSize(con->destination, con->origin, nodes);
        con->disabled = true;
        res += a * b * con->speed;
    }
    return res;
}
 
int main() {
    long n;
    cin >> n;
    
    vector<Connection*> nodes [n + 1];
    vector<Connection*> connections;
 
    for (long i = 0; i < n - 1; i++) {
        long a, b, x;
        cin >> a >> b >> x;
        Connection* con = new Connection(x, a, b);
        nodes[a].push_back(con);
        nodes[b].push_back(con);
        connections.push_back(con);
    }
 
    long output = solve(nodes, connections);
    
    cout << output;
}

Test 1

Group: 1, 2, 3

Verdict: ACCEPTED

Test 2

Group: 2, 3

Verdict: ACCEPTED

Test 3

Group: 3

Verdict: TIME LIMIT EXCEEDED