Tietoverkko

#include <iostream>
#include <vector>
#include <utility>
#include <unordered_map>
using namespace std;

static vector<vector<pair<uint32_t, uint64_t>>> puu;

static unordered_map<uint64_t, uint64_t> muistio;

static uint64_t hitaammat(uint32_t s, uint32_t e, uint64_t vert) {
    vector<pair<uint32_t, uint32_t>> stack;
    stack.push_back({s, e});
    uint64_t summa = 0;
    while (!stack.empty()) {
        auto newval = stack.back();
        s = newval.first;
        e = newval.second;
        stack.pop_back();
        for (auto& seur : puu[s]) {
                if (seur.first == e)
                    continue;
                if (seur.second < vert)
                    continue;
                uint64_t key = (uint64_t)seur.first << 32 | (uint64_t)s;
                auto muistioviite = muistio.find(key);
                if (muistioviite == muistio.end() || seur.second != vert + 1) {
                    ++summa;
                    stack.push_back({seur.first, s});
                } else {
                    summa += muistioviite->second + 1;
                }
        }
    }
    muistio[(uint64_t)s << 32 | (uint64_t)e] = summa;
    return summa;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    uint64_t n;
    cin >> n;

    puu.reserve(n);

    unordered_map<uint32_t, uint32_t> count;

    for (uint32_t i = 0; i < n - 1; ++i) {
        uint32_t a, b;
        uint64_t nopeus;
        cin >> a >> b >> nopeus;
        puu[a - 1].push_back({b - 1, 200001 * nopeus + count[n]});
        puu[b - 1].push_back({a - 1, 200001 * nopeus + count[n]});
        ++count[n];
    }

    uint64_t summa = 0;
    for (uint32_t i = 0; i < n; ++i) {
        for (auto& solmu : puu[i]) {
            if (i > solmu.first)
                continue;
            summa += (hitaammat(solmu.first, i, solmu.second) + 1) * (hitaammat(i, solmu.first, solmu.second) + 1) * (solmu.second / 200001);
        }
    }

    cout << summa << '\n';
}

Test 1

Group: 1, 2, 3

Verdict: ACCEPTED

Test 2

Group: 2, 3

Verdict: ACCEPTED

Test 3

Group: 3

Verdict: TIME LIMIT EXCEEDED