Posting my solution from last week (spiffed up a bit). It passes the tests, but is pretty slow. Gonna try to improve it this week using a new technique...
class Solution:
def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
## convert the given edges into a more useful structure
neighbors_by_node = defaultdict(list)
for edge in edges:
neighbors_by_node[edge[0]].append(edge[1])
neighbors_by_node[edge[1]].append(edge[0])
## perform a DFS starting at node 1, recording each node's
## parent and distance from the root
distance_to_node_1_by_node = {1: 0}
parent_by_node = {1: None}
def calc_distances_and_parents (node, node_parent):
for neighbor in neighbors_by_node[node]:
if neighbor != node_parent:
distance_to_node_1_by_node[neighbor] = 1 + distance_to_node_1_by_node[node]
parent_by_node[neighbor] = node
calc_distances_and_parents(neighbor, node)
calc_distances_and_parents(1, None)
## a memoized function for finding the distance between two nodes
distance_memo = {}
def find_distance_between(node_a, node_b):
## put the nodes into a canonical ordering to avoid
## wasted space / missed lookups
key = (node_a,node_b) if node_a<node_b else (node_b,node_a)
if key not in distance_memo:
if node_a == node_b:
result = 0
elif node_a == 1:
result = distance_to_node_1_by_node[node_b]
elif node_b == 1:
result = distance_to_node_1_by_node[node_a]
else:
if distance_to_node_1_by_node[node_a] > distance_to_node_1_by_node[node_b]:
result = 1 + find_distance_between(parent_by_node[node_a], node_b)
else:
result = 1 + find_distance_between(parent_by_node[node_b], node_a)
distance_memo[key] = result
return distance_memo[key]
## a memoized calculation for the number of solutions
## given the path length
@cache
def find_num_ways(N_edges):
if N_edges == 0:
return 0
else:
return pow(2, N_edges-1, mod=10**9 + 7)
## process each query and update the result
result = []
for q in queries:
N_edges = find_distance_between(q[0],q[1])
N_ways = find_num_ways(N_edges)
result.append(N_ways)
return result