Medium
Tree Diameter — Python
Full explanation · Time O(\ · Space V\
# Time: O(|V| + |E|)
# Space: O(|E|)
# iterative dfs
class Solution(object):
def treeDiameter(self, edges):
"""
:type edges: List[List[int]]
:rtype: int
"""
def iter_dfs():
result = 0
stk = [(1, (0, -1, [0]))]
while stk:
step, args = stk.pop()
if step == 1:
u, p, ret = args
for v in reversed(adj[u]):
if v == p:
continue
ret2 = [0]
stk.append((2, (ret2, ret)))
stk.append((1, (v, u, ret2)))
elif step == 2:
ret2, ret = args
result = max(result, ret[0]+(ret2[0]+1))
ret[0] = max(ret[0], ret2[0]+1)
return result
adj = [[] for _ in range(len(edges)+1)]
for u, v in edges:
adj[u].append(v)
adj[v].append(u)
return iter_dfs()
# Time: O(|V| + |E|)
# Space: O(|E|)
# dfs
class Solution2(object):
def treeDiameter(self, edges):
"""
:type edges: List[List[int]]
:rtype: int
"""
def dfs(u, p):
mx = 0
for v in adj[u]:
if v == p:
continue
curr = dfs(v, u)
result[0] = max(result[0], mx+(curr+1))
mx = max(mx, curr+1)
return mx
adj = [[] for _ in range(len(edges)+1)]
for u, v in edges:
adj[u].append(v)
adj[v].append(u)
result = [0]
dfs(0, -1)
return result[0]
# Time: O(|V| + |E|)
# Space: O(|E|)
# bfs, tree dp
class Solution3(object):
def treeDiameter(self, edges):
"""
:type edges: List[List[int]]
:rtype: int
"""
def bfs():
result = 0
dp = [0]*len(adj)
degree = map(len, adj)
q = [u for u in xrange(len(degree)) if degree[u] == 1]
while q:
new_q = []
for u in q:
if degree[u] == 0:
continue
degree[u] -= 1
for v in adj[u]:
if degree[v] == 0:
continue
result = max(result, dp[v]+(dp[u]+1))
dp[v] = max(dp[v], (dp[u]+1))
degree[v] -= 1
if degree[v] == 1:
new_q.append(v)
q = new_q
return result
adj = [[] for _ in range(len(edges)+1)]
for u, v in edges:
adj[u].append(v)
adj[v].append(u)
return bfs()
# Time: O(|V| + |E|)
# Space: O(|E|)
# bfs
class Solution4(object):
def treeDiameter(self, edges):
"""
:type edges: List[List[int]]
:rtype: int
"""
def bfs(root):
d = new_root = -1
lookup = [False]*len(adj)
lookup[root] = True
q = [root]
while q:
d, new_root = d+1, q[0]
new_q = []
for u in q:
for v in adj[u]:
if lookup[v]:
continue
lookup[v] = True
new_q.append(v)
q = new_q
return d, new_root
adj = [[] for _ in range(len(edges)+1)]
for u, v in edges:
adj[u].append(v)
adj[v].append(u)
_, root = bfs(0)
d, _ = bfs(root)
return d