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 > 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 union
ing 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.