1974 - TrumpLandia

După o confruntare eroică, Trump a fost ales democratic președintele unei republici prospere, TrumpLandia. Natural, prima sa prioritate e să construiască o rețea de buncăre, pentru a veni în întâmpinarea unui atac al vecinilor imperaliști.

N buncăre au fost deja construite. Trump trebuie să aleagă între M posibile rute bidirecționale pentru a le conecta. Un plan de conectare este alcătuit dintr-un subset minim de astfel de legături, astfel încât din fiecare locație să se poată ajunge în orice altă locație. Costul unui plan este dat de suma distanțelor legăturilor.

Deoarece rutele sunt vulnerabile la bombardamente aeriene, Trump vă cere să estimați un plan de cost minim. În plus, Trump vă cere să calculați Q soluții de rezervă, pentru situația în care una dintre cele M legături este sigur inclusă în plan.

Date de intrare

Datele de intrare vor fi citite din fișierul trumplandia.in. Pe prima linie se află trei numere separate prin spațiu: N (numărul de buncăre), M (numărul de rute posibile), Q (numărul de soluții de rezervă).

Pe următoarele M linii se află triplete de forma (nod1, nod2, dist12) ce descriu o posibilă legătură între nodurile nod1 și nod2 cu distanța dist12.

Pe următoarele Q linii se află câte un număr Qi, reprezentând indexul unei muchii care trebuie inclusă într-un plan de rezervă. (valorile se pot repeta)

Date de ieșire

Datele de ieșire vor fi afișate în fișierul trumplandia.out.

Pe prima linie se va afla un număr reprezentând costul celui mai bun plan, pentru situația în care pot fi folosite oricare dintre cele M rute.

Pe următoarele Q linii se află un număr Ci reprezentând costul minim al unui plan care conține muchia Qi.

Restricții și precizări

• 1 ≤ N ≤ 200.000 (pentru 30% din teste N ≤ 1000)
• N-1 ≤ M ≤ 200.000 (pentru 30% din teste M ≤ 1000)
• 1 ≤ Q ≤ 300.000 (pentru 30% din teste Q ≤ 1000)
• distanța asociată unei rute este o valoare întreagă din intervalul [1, 1.000.000.000].

ATENȚIE!!! Între oricare două buncăre pot exista MAI MULTE rute de distanțe diferite sau nu. Exemplu: (1 2 4) și (1 2 6) sunt două muchii între buncărele 1 și 2 de distanțe 4, respectiv 6.

Exemplu:

trumplandia.in

6 8 8
1 2 4
4 1 6
3 1 2
2 3 3
5 2 4
3 4 3
4 5 5
5 6 1
1
2
3
4
5
6
7
8

trumplandia.out

13
14
16
13
13
13
13
14
13

Rezolvare

with open("trumplandia.in", "r") as fin:
    N, M, Q = map(int, fin.readline().split())
    edges = [tuple(map(int, fin.readline().split())) for _ in range(M)]
    reserve_edges = [int(fin.readline()) for _ in range(Q)]




def find(parent, node):
    if parent[node] == node:
        return node
    parent[node] = find(parent, parent[node])
    return parent[node]


def union(parent, rank, x, y):
    x = find(parent, x)
    y = find(parent, y)
    if x == y:
        return False
    if rank[x] < rank[y]:
        parent[x] = y
    elif rank[x] > rank[y]:
        parent[y] = x
    else:
        parent[y] = x
        rank[x] += 1
    return True


def kruskal(edges, N):
    edges.sort(key=lambda x: x[2])  

    parent = list(range(N + 1))  
    rank = [0] * (N + 1) 
    min_cost = 0  

    mst_edges = []  

    for edge in edges:
        u, v, cost = edge
        if find(parent, u) != find(parent, v):
            union(parent, rank, u, v)
            min_cost += cost
            mst_edges.append(edge)

    return min_cost, mst_edges


min_cost, mst_edges = kruskal(edges, N)


adj_list = {i: [] for i in range(1, N + 1)}
for edge in mst_edges:
    u, v, cost = edge
    adj_list[u].append((v, cost))
    adj_list[v].append((u, cost))


with open("trumplandia.out", "w") as fout:
    fout.write(str(min_cost) + '\n')
    for reserve_edge in reserve_edges:
        reserve_cost = min_cost  

        if reserve_edge not in [edge[:2] for edge in mst_edges]:
            u, v, cost = edges[reserve_edge - 1]
            for edge in mst_edges:
                if edge[0] == u or edge[1] == u or edge[0] == v or edge[1] == v:
                    reserve_cost += cost - edge[2]
                    break

        fout.write(str(reserve_cost) + '\n')