CSES - Datatähti 2022 alku - Results
Submission details
Task:Tietoverkko
Sender:tonero
Submission time:2021-10-16 16:07:11 +0300
Language:C++17
Status:READY
Result:10
Feedback
groupverdictscore
#1ACCEPTED10
#20
#30
Test results
testverdicttimegroup
#1ACCEPTED0.01 s1, 2, 3details
#20.02 s2, 3details
#3--3details

Code

#include <bits/stdc++.h>
#pragma GCC optimize("Ofast")
#pragma GCC target("avx,avx2,fma")

using namespace std;

#define ll long long
#define ull unsigned long long

struct pair_hash {
    inline ll operator()(const std::pair<int,int> & v) const {
        return ((ll)v.first) << 32 | (ll)(unsigned int)v.second;
    }
};

vector<pair<int, int>>* vecs[200001] = {};
//map of integer, amount
std::unordered_map<int, vector<unordered_map<int,int>>> spds;

ull cur = 0;

//TODO: id 2, big tree

void search_tree(int start, int spd, int lst){

    vector<pair<unordered_map<int,int>, int>> nspds = {};

    unordered_map<int,int> loneSpds = {};

    for (const auto &item : *vecs[start]){
        if(item.first == lst) continue;
        if(vecs[item.first]->size() > 1){
            search_tree(item.first, item.second, start);
            vector<unordered_map<int,int>> spdVec = spds[item.first];
            for (const auto &item2 : spdVec){
                nspds.emplace_back(item2, item.first);
            }
        }else{
            const auto iter = loneSpds.find(item.second);
            int val = 1;
            if(iter != loneSpds.end()) {
                val += iter->second;
                iter->second = val;
            } else
                loneSpds.emplace(item.second, val);
        }

    }

    for (const auto &item : nspds){
        for (const auto &item2 : item.first){
            for (const auto &item3 : loneSpds){
                ull tmp = min(item3.first, item2.first);
                cur += tmp*(item3.second * item2.second);
            }
        }
    }

    for (const auto &item : *vecs[start]){
        if(item.first == lst || vecs[item.first]->size() == 1) continue;
        for (const auto &item2 : nspds){
            if(item2.second == item.first){
                continue;
            }
            for (const auto &item3 : item2.first){
                ull tmp = min(item.second, item3.first);
                cur += tmp*(item3.second);
            }
        }
        const auto iter = loneSpds.find(item.second);
        int val = 1;
        if(iter != loneSpds.end()) {
            val += iter->second;
            iter->second = val;
        } else
            loneSpds.emplace(item.second, val);
    }


    unordered_map<int,int> ret_spds = {};
    //test
    auto loopNspds = nspds;
    for (auto iter = loopNspds.begin(); iter != loopNspds.end(); ) {

        for (const auto &nspd2 : loopNspds){
            if(&*iter == &nspd2) continue;
            for (const auto &item : iter->first){
                for (const auto &item2 : nspd2.first){
                    ull tmp = min(item.first, item2.first);
                    cur += tmp*(item.second * item2.second);
                }
            }
        }

        for (const auto &item : iter->first){
            int retSpd = min(spd, item.first);
            const auto iter2 = ret_spds.find(retSpd);
            int val = item.second;
            if(iter2 != ret_spds.end()) {
                val += iter2->second;
                iter2->second = val;
            } else
                ret_spds.emplace(retSpd, val);
        }
        iter = loopNspds.erase(iter);
    }

    for (const auto &item : nspds){
        for (const auto &item2 : item.first){
            cur += item2.first * item2.second;
        }
    }

    //add all spds to loneSpds and spds between loneSpds in the same pair
    unordered_map<int, int> nLoneSpds = {};
    for (const auto &item : loneSpds){
        cur += item.first*item.second;
        int am = item.second -1;
        cur += item.first*am;
    }
    //
    for (auto iter = loneSpds.begin(); iter != loneSpds.end(); ) {
        auto item = *iter;
        for (const auto &item2 : loneSpds){
            if(&item2 == &*iter) continue;
            ull tmp = min(item.first, item2.first);
            //mby + ?
            cur += tmp*((item.second) * item2.second);
        }

        if(start != 1){
            //push loneSpds to nSpds
            int nspd = min(item.first, spd);
            const auto iter2 = nLoneSpds.find(nspd);
            int amount = item.second;
            if(iter2 != nLoneSpds.end()){
                amount += iter2->second;
                iter2->second = amount;
            }else{
                nLoneSpds.emplace(nspd, amount);
            }
        }

        iter = loneSpds.erase(iter);
    }

    for (const auto &item : nLoneSpds){
        const auto iter = ret_spds.find(item.first);
        int val = item.second;
        if(iter != ret_spds.end()) {
            val += iter->second;
            iter->second = val;
        } else
            ret_spds.emplace(item.first, val);
    }

    vector<unordered_map<int,int>> retSpdVec = {ret_spds};

    spds[start] = retSpdVec;
}

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

    int n;
    cin >> n;

    for (int i = 1; i < n; i++)
    {
        int a,b,x;
        cin >> a >> b >> x;

        if(vecs[a] == nullptr) {
            vecs[a] = new vector<pair<int, int>>();
        }
        if(vecs[b] == nullptr) {
            vecs[b] = new vector<pair<int, int>>();
        }

        vecs[a]->emplace_back(make_pair(b,x));
        vecs[b]->emplace_back(make_pair(a,x));
    }

    search_tree(1, __INT32_MAX__, -1);

    cout << cur;
    flush(std::cout);
}

Test details

Test 1

Group: 1, 2, 3

Verdict: ACCEPTED

input
100
1 2 74
1 3 100
2 4 50
3 5 40
...

correct output
88687

user output
88687

Test 2

Group: 2, 3

Verdict:

input
5000
1 2 613084013
1 3 832364259
2 4 411999902
3 5 989696303
...

correct output
1103702320243776

user output
1100219101766720

Test 3

Group: 3

Verdict:

input
200000
1 2 613084013
1 3 832364259
2 4 411999902
3 5 989696303
...

correct output
1080549209850010931

user output
(empty)