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.head = None
        self.tail = None
        self.count = 0

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

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

    def add_node(self, node):
        if not self.rep.head:
            self.rep.head = node
        self.rep.count += 1
        self.rep.tail = node
        node.head = self.rep

def make_set(v):
    new_linked_list = DisjointSetLinkedList()
    new_head_obj = DisjointSetNode(v)
    new_linked_list.add_node(new_head_obj)
    return new_linked_list, new_head_obj

def find_set(v, node_map):
    return node_map[v].head

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

    if u_node.head.count &gt; v_node.head.count:
        small_rep = v_node.head
        large_rep = u_node.head
    else:
        small_rep = u_node.head
        large_rep = v_node.head

    # update all nodes to point to new rep
    temp = small_rep.head
    while temp:
        temp.head = large_rep
        temp = temp.next
    # update last node in first list to point to head of second list
    large_rep.tail.next = small_rep.head
    # 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)
        sets_map[new_node.head] = new_set
        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.
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 &gt;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.