Disjoint Set Connected Components With Weighted Graph

I have been trying to solve this HackerRank problem (link).

The basic premise of this problem is that there is a tree with undirected, but weighted, edges. The cost of a path in this tree is taken to be the MAXIMUM cost of any edge in the path. I will be given a series of queries in the form of [L, R] and I have to output how many paths in that tree have a cost in the provided inclusive range.

This is the code I've written -

class DisjointSetRep():
def __init__(self):
self.tail = None
self.count = 0

class DisjointSetNode():
def __init__(self, val):
self.val = val
self.next = None

def __init__(self):
self.rep = DisjointSetRep()

self.rep.count += 1
self.rep.tail = node

def make_set(v):

def find_set(v, node_map):

def union(u, v, node_map, sets_map):
u_node, v_node = node_map[u], node_map[v]

else:

# update all nodes to point to new rep
while temp:
temp = temp.next
# update last node in first list to point to head of second list
# update new tail
large_rep.tail = small_rep.tail
# update count
large_rep.count += small_rep.count
del sets_map[small_rep]
return large_rep

def create_data(edges, node_map, sets_map):

# sort the edges first, according to cost
edges.sort(key=lambda x:x[2])
cost_map = {} # key - cost, value - no of paths
largest_cost = edges[-1][2]

for edge in edges:
if find_set(edge[0], node_map) != find_set(edge[1], node_map):
unioned_set = union(edge[0], edge[1], node_map, sets_map)
cost_map[edge[2]] = unioned_set.count - 1

prefixed_cost_data = {0: 0}
for i in xrange(1, largest_cost+1):
val = cost_map.get(i, 0)
prefixed_cost_data[i] = prefixed_cost_data[i-1] + val

return prefixed_cost_data

if __name__ == '__main__':
n, q = map(int, raw_input().split())
sets_map = {}
node_map = {}
for i in xrange(1, n+1):
new_set, new_node = make_set(i)
node_map[i] = new_node
edges = []
for _ in xrange(n-1):
edges.append(map(int, raw_input().split()))
prefixed_cost_data = create_data(edges, node_map, sets_map)
for _ in xrange(q):
l, r = map(int, raw_input().split())
print prefixed_cost_data[r] - prefixed_cost_data[l-1]


Let me explain the logic above, which I have derived from this comment -

I sort the edges according to their cost. I then iterate over them and construct the tree edge-by-edge by unioning each vertex, which is initially a disjoint set containing itself (make_set). At any point, no_of_vertexes - 1 gives the no of paths in the tree that contain the maximum cost, which is what I use in unioned_set.count - 1.

This gives me a cost_map with keys as the costs and the values as the number of paths. I also generated a prefixed sum array so that to get the output for [L, R], instead of calculating the no of paths for each value in the range, I can just do prefixed_cost_data[r] - prefixed_cost_data[l-1].

The implementation of the disjoint set is taken straight from CLRS (Section 21.2).

I think the above logic is correct, but I guess it's too slow since most of the test cases timeout.

Can anyone help me in optimizing it? I guess the entire logic needs to be revamped.

• I don't understand why you get a timeout, but I think your path counting is wrong: When you add the next-heaviest edge, the number of paths having that weight is equal to the product of the number of vertices in each of the components it connects (in the simple case where there is a unique edge of this weight; when multiple equal-weight edges connect a larger subtree of components into a single component, the calculation is more complicated). – j_random_hacker Aug 11 at 15:28

Using arrays instead of classes and objects would make it easier for you to implement the disjoint set union concept.

Here's a resource to start off with:- Disjoint Set Union

for edge in edges:
if find_set(edge[0], node_map) != find_set(edge[1], node_map):
unioned_set = union(edge[0], edge[1], node_map, sets_map)
cost_map[edge[2]] = unioned_set.count - 1


cost_map[edge[2]] will not be equal to the size of resultant set after union - 1.

It will be equal to the product of size of two individual sets.

Over here,

for i in xrange(1, largest_cost+1):
val = cost_map.get(i, 0)
prefixed_cost_data[i] = prefixed_cost_data[i-1] + val


largest_cost can be >10^9. Neither can you store an array of that size, nor can you execute a loop with 10^9 operations with a time limit of 1sec.

Try using binary search on the L R query.

If you implement disjoint union set with arrays and make the above corrections , then I believe the code should work.