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)
    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 > v_node.head.count:
        small_rep = v_node.head
        large_rep = u_node.head
        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.

  • $\begingroup$ 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). $\endgroup$ – j_random_hacker Aug 11 at 15:28

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.