Submission #758970

#	Time	Username	Problem	Language	Result	Execution time	Memory
758970		thimote75	Toll (APIO13_toll)	C++14	0 / 100	1 ms	212 KiB

toll

#include <bits/stdc++.h>

using namespace std;

#define int long long
using idata = vector<int>;
using bdata = vector<bool>;

struct Road {
    int cost;
    int a, b;

    Road (int cost, int a, int b) 
    : cost(cost), a(a), b(b) {}

    int next (int node) {
        if (node == a) return b;
        return a;
    }
};
struct TaxRoad {
    int next;
    bool tax;

    TaxRoad (int n, bool t) : next(n), tax(t) {}
};

using road_array = vector<Road*>;

struct DSU {
private:
    idata parents;
public:
    DSU (int size) {
        parents.resize(size);

        for (int i = 0; i < size; i ++)
            parents[i] = i;
    }
    void reset () {
        for (int i = 0; i < parents.size(); i ++)
            parents[i] = i;
    }

    int boss (int x) {
        if (parents[x] != x)
            parents[x]  = boss(parents[x]);
        
        return parents[x];
    }
    bool merge (int a, int b) {
        a = boss(a);
        b = boss(b);

        parents[a] = b;

        return a != b;
    }
};

road_array compute_mst (int N, road_array roads) {
    DSU dsu (N);

    road_array mst;

    for (Road* road : roads) {
        if (!dsu.merge(road->a, road->b)) continue ;

        mst.push_back(road);
    }

    return mst;
}

idata compress_metadata (int N, road_array &current_road_array, set<int> kernel, idata &people_count) {
    idata result(N);

    vector<set<Road*>> roads(N);

    idata degree(N);
    for (Road* road : current_road_array) {
        degree[road->a] ++;
        degree[road->b] ++;

        roads[road->a].insert(road);
        roads[road->b].insert(road);
    }

    set<pair<int, int>> degree_queue;
    for (int i = 0; i < N; i ++) {
        if (kernel.find(i) != kernel.end()) continue ;

        degree_queue.insert({ degree[i], i });
    }

    while (degree_queue.size() != 0) {
        auto h = *(degree_queue.begin());
        degree_queue.erase(h);

        int d = h.first; int node = h.second;
        if (d >= 3) break ;

        result[node] = -1;

        if (d == 1) {
            Road* road = *(roads[node].begin());
            int   next = road->next(node);

            roads[next].erase(road);
            roads[node].erase(road);

            people_count[next] += people_count[node];

            if (degree_queue.find({ degree[next], next }) != degree_queue.end())
                degree_queue.erase({ degree[next], next });
            degree[next] --;
            if (kernel.find(next) == kernel.end())
                degree_queue.insert({ degree[next], next });
        } else if (d == 2) {
            Road* road1 = *(roads[node].begin());
            int   next1 = road1->next(node);

            roads[node ].erase(road1);
            roads[next1].erase(road1);

            Road* road2 = *(roads[node].begin());
            int   next2 = road2->next(node);

            roads[node ].erase(road2);
            roads[next2].erase(road2);
        
            Road* substitute = new Road(max(road1->cost, road2->cost), next1, next2);

            if (road1->cost > road2->cost) people_count[next2] += people_count[node];
            else people_count[next1] += people_count[node];
        
            roads[next1].insert(substitute);
            roads[next2].insert(substitute);
        }
    }

    int new_index = 0;
    for (int i = 0; i < N; i ++)
        if (result[i] != -1)
            result[i] = new_index ++;

    current_road_array.clear();

    for (int i = 0; i < N; i ++) {
        for (Road* road : roads[i]) {
            if (road->next(i) < i) continue ;
            int next = road->next(i);

            road->a = result[i];
            road->b = result[next];
            current_road_array.push_back(road);
        }
    }

    return result;
}

vector<vector<TaxRoad>> tax_graph;

idata greedy_roads_weight;
idata parents;
idata depth;

void lca_dfs (int node, int parent, int _depth) {
    depth[node] = _depth;
    parents[node] = parent;

    for (TaxRoad &next : tax_graph[node]) if (next.next != parent)
        lca_dfs(next.next, node, _depth + 1);
}

/**
 * Calcul des poids associés à une route pour Mr. Greedy
 * 
 * Est égal au minimum des valeurs pour les routes ne faisant pas parties du MST pour lesquels le chemin dans le MST passe par cette route.
 * 
 * On peut alors calculer en chaque noeuds ces valeurs avec un algorithme semblable à celui du lca en O(K)
 */
void update_weight (Road* out_of_mst) {
    int a = out_of_mst->a;
    int b = out_of_mst->b;
    if (depth[a] > depth[b]) swap(a, b);

    while (depth[a] != depth[b]) {
        greedy_roads_weight[b] = min(
            greedy_roads_weight[b],
            out_of_mst->cost
        );

        b = parents[b];
    }
    while (a != b) {
        greedy_roads_weight[a] = min(
            greedy_roads_weight[a],
            out_of_mst->cost
        );
        greedy_roads_weight[b] = min(
            greedy_roads_weight[b],
            out_of_mst->cost
        );

        a = parents[a];
        b = parents[b];
    }
}

pair<int, int> res_dfs (idata &people_count, int node, int parent) {
    pair<int, int> result = { people_count[node], 0 };
 
    for (TaxRoad &road : tax_graph[node]) if (road.next != parent) {
        pair<int, int> local = res_dfs(people_count, road.next, node);
 
        result.first  += local.first;
        result.second += local.second;
 
        if (!road.tax) continue ;

        result.second += greedy_roads_weight[road.next] * local.first;
    }

    return result;
}

signed main () {
    int N, M, K;
    cin >> N >> M >> K;

    road_array road_array;

    for (int i = 0; i < M; i ++) {
        int a, b, c;
        cin >> a >> b >> c;
        a --; b --;

        Road* road = new Road(c, a, b);

        road_array.push_back(road);
    }

    road_array = compute_mst(N, road_array);

    set<int> kernel;
    kernel.insert(0);

    vector<pair<int, int>> greedy_roads;

    for (int i = 0; i < K; i ++) {
        int a, b;
        cin >> a >> b;
        a --; b --;

        kernel.insert(a);
        kernel.insert(b);

        greedy_roads.push_back({ a, b });
    }

    idata people_count(N);
    for (int i = 0; i < N; i ++)
        cin >> people_count[i];

    idata compressor = compress_metadata(N, road_array, kernel, people_count);

    for (auto &x : greedy_roads) {
        x.first  = compressor[x.first];
        x.second = compressor[x.second];
    }

    int nC = 0;
    for (int u : compressor) if (u != -1) nC ++;

    // TODO iterate over all subsets.
    DSU dsu(nC);

    idata compressed_pc(nC);
    for (int i = 0; i < N; i ++) {
        if (compressor[i] == -1) continue ;

        compressed_pc[compressor[i]] = people_count[i];
    }

    tax_graph.resize(nC);
    depth.resize(nC);
    parents.resize(nC);

    int result = 0;

    for (int mask = 0; mask < (1 << K); mask ++) {
        bool has_cycle = false;
        greedy_roads_weight.clear();
        greedy_roads_weight.resize(nC, 1e9);

        for (int i = 0; i < nC; i ++) tax_graph[i].clear();

        dsu.reset();
        for (int iK = 0; iK < K; iK ++) {
            if (((1 << iK) & mask) == 0) continue ;

            if (!dsu.merge( greedy_roads[iK].first, greedy_roads[iK].second )) {
                has_cycle = true;
                break;
            }

            tax_graph[greedy_roads[iK].first ].push_back(TaxRoad( greedy_roads[iK].second, true ));
            tax_graph[greedy_roads[iK].second].push_back(TaxRoad( greedy_roads[iK].first,  true ));
        }

        if (has_cycle) continue ;

        vector<Road*> out_of_mst;

        for (Road* road : road_array) {
            if (!dsu.merge(road->a, road->b)) {
                out_of_mst.push_back(road);
                continue ;
            }

            tax_graph[road->a].push_back(TaxRoad(road->b, false));
            tax_graph[road->b].push_back(TaxRoad(road->a, false));
        }

        lca_dfs(0, -1, 0);

        for (Road* road : out_of_mst)
            update_weight(road);

        int answer = res_dfs(compressed_pc, 0, -1).second;

        result = max(answer, result);
    }

    cout << result << endl;
}

Compilation message (stderr)

toll.cpp: In member function 'void DSU::reset()':
toll.cpp:42:27: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   42 |         for (int i = 0; i < parents.size(); i ++)
      |                         ~~^~~~~~~~~~~~~~~~

• Subtask 1: 0 / 16
• Subtask 2: 0 / 18
• Subtask 3: 0 / 22
• Subtask 4: 0 / 22
• Subtask 5: 0 / 22