PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Simple
PREREQUISITES:
Sorting, understanding Kruskal’s MST algorithm.
PROBLEM:
There’s a complete graph on N vertices.
You know the weights of its edges (given as an array W), but not which weight corresponds to which edge.
Across all possible ways of assigning the weights to edges, find the maximum possible weight of the MST of the graph.
EXPLANATION:
To solve this problem, it’s useful to understand how Kruskal’s algorithm to compute the MST works.
Kruskal’s algorithm is as follows:
- Let E be the set of edges, where each edge is a tuple (u, v, w) denoting an edge of weight w between u and v.
- The elements of E are processed in ascending order of weight.
- When processing the edge (u, v, w): if u and v are already in the same connected component when considering only previously chosen edges, do nothing - i.e. skip this edge.
Otherwise, add w to the MST weight, and include edge (u, v) into the MST.
It can be proved that this algorithm will output an MST always.
Let’s try to apply the idea behind this algorithm to our problem.
Our goal is to maximize the sum of chosen weights.
Since Kruskal processes edges in increasing order of weight, this means we want to skip smaller weights as much as possible - and an edge can be skipped only when its endpoints are already connected.
Further, if we’ve chosen K edges in the MST already, we can only skip a total of \frac{K\cdot (K+1)}{2} - K = \frac{K\cdot (K-1)}{2} other edges before we’re forced to choose one.
This is because we can only skip edges between vertices that are in the same connected component; and with K edges the best we can do is if they all form a single component (with K+1 vertices).
This naturally gives rise to a simple greedy algorithm, where we process edges in ascending order of weight, and choose one only when we’re forced to.
That is, let’s first sort W in ascending order, so that W_1\le W_2\le W_3\le\ldots
Now,
- We initially have 0 edges.
This allows us to skip 0 edges, so we’re forced to choose the first edge, with weight W_1. - We now have 1 edges.
This allows us to skip \frac{1\cdot 0}{2} = 0 edges, so we’re forced to choose the next edge, with weight W_2. - Now we have 2 edges.
This allows us to skip \frac{2\cdot 1}{2} = 1 edge, so we skip the next edge and choose W_4 instead. - Next, we have three edges.
This allows us to skip 3 edges in total.
However, we’ve already skipped one edge, so we can only skip two more - meaning we skip W_5 and W_6, and choose W_7. - Next, with four edges, we’re allowed to skip 6 edges in total.
Three of them have already been skipped; so we skip just the next three and pick W_{11} instead.
This algorithm continues till we’ve picked N-1 edges.
In general, after picking the K-th edge, we can skip the next K-1 edges.
This leads to the picked edges being:
or more generally, all edges with index \frac{K\cdot (K+1)}{2} + 1 for some K \ge 0 (after sorting W).
The above intuitive greedy algorithm is indeed correct, below is a formal proof of why.
Proof
First, we need to show the weights we chose are actually achievable as an MST.
This is not too hard to do, by virtue of how we chose them.
Specifically, for each K \ge 0, do the following:
- Let x = \frac{K\cdot (K+1)}{2} + 1 be the index of the weight we’re choosing.
- Add an edge of weight W_x between vertices 1 and K+2.
- Then, add edges of weights W_{x+1}, W_{x+2}, \ldots, W_{x+K} between K+2 and 2, 3, 4, \ldots, K+1.
This construction ensures that the edges chosen for the MST are exactly the ones with weights our algorithm picked; those being edges (1, 2), (1, 3), (1, 4), \ldots, (1, N).
Essentially, the idea is to connect each new vertex to 1 using an edge that’s going to be chosen, and then connect the new vertex to all previous non-1 vertices using edges that are going to be skipped.We also need to show that no larger choice works, but that’s easy enough to see by very choice of the greedy - we skipped as many low-weight edges as we possibly could till we were forced to choose.
If you’d like to write out a formal proof, it’s possible to use exchange arguments to show this.
TIME COMPLEXITY:
\mathcal{O}(N^2 \log N) per testcase.
CODE:
Editorialist's code (PyPy3)
for _ in range(int(input())):
n = int(input())
w = list(map(int, input().split()))
w.sort()
ans, skip, v = 0, 0, 0
for wt in w:
if skip == 0:
ans += wt
skip = v
v += 1
else: skip -= 1
print(ans)