Task · algorithms hard Python ·

Heaviest edge on tree paths

def path_max(n: int, edges: list[list[int]], queries: list[list[int]]) -> list[int]:

path_max(4, [[0, 1, 5], [1, 2, 3], [1, 3, 8]], [[0, 2], [2, 3], [3, 3]])  # -> [5, 8, 0]
path_max(2, [[0, 1, 7]], [[1, 0]])  # -> [7]

• #1MiniMax M3.1 Flashsolved · 9 of 9 testsanswer in 4.6 s · runtime 1.0 s · $0.0010

  ✓	path_max(4, [[0, 1, 5], [1, 2, 3], [1, 3, 8]], [[0, 2], [2, 3], [3, 3]])expected [5, 8, 0]	<1 ms
  ✓	path_max(2, [[0, 1, 7]], [[1, 0]])expected [7]	<1 ms

  ✓ 2 ms✓ <1 ms✓ <1 ms✓ <1 ms✓ <1 ms⏱ 482 ms⏱ 550 ms
  Code

  def path_max(n: int, edges: list[list[int]], queries: list[list[int]]) -> list[int]:
      if n <= 1:
          return [0] * len(queries)
  
      adj = [[] for _ in range(n)]
      for e in edges:
          u = e[0]; v = e[1]; w = e[2]
          adj[u].append((v, w))
          adj[v].append((u, w))
  
      LOG = n.bit_length()
      parent = [-1] * n
      depth = [0] * n
      parent_w = [0] * n
  
      parent[0] = 0
      stack = [0]
      while stack:
          u = stack.pop()
          pu = parent[u]
          du = depth[u] + 1
          for v, w in adj[u]:
              if v == pu:
                  continue
              parent[v] = u
              depth[v] = du
              parent_w[v] = w
              stack.append(v)
  
      up = [parent]
      mx = [parent_w]
      for k in range(1, LOG):
          pu = up[k - 1]
          pmx = mx[k - 1]
          nu = [0] * n
          nmx = [0] * n
          for v in range(n):
              a = pu[v]
              nu[v] = pu[a]
              x = pmx[v]
              y = pmx[a]
              nmx[v] = x if x > y else y
          up.append(nu)
          mx.append(nmx)
  
      res = []
      append = res.append
      for q in queries:
          a = q[0]; b = q[1]
          if a == b:
              append(0)
              continue
          best = 0
          if depth[a] < depth[b]:
              a, b = b, a
          d = depth[a] - depth[b]
          k = 0
          while d:
              if d & 1:
                  w = mx[k][a]
                  if w > best:
                      best = w
                  a = up[k][a]
              d >>= 1
              k += 1
          if a != b:
              for k in range(LOG - 1, -1, -1):
                  ua = up[k][a]
                  ub = up[k][b]
                  if ua != ub:
                      w = mx[k][a]
                      if w > best:
                          best = w
                      w = mx[k][b]
                      if w > best:
                          best = w
                      a = ua
                      b = ub
              w = mx[0][a]
              if w > best:
                  best = w
              w = mx[0][b]
              if w > best:
                  best = w
          append(best)
      return res

• #1DeepSeek V4.1 Flashsolved · 9 of 9 testsanswer in 5.8 s · runtime 1.0 s · $0.0005

  ✓	path_max(4, [[0, 1, 5], [1, 2, 3], [1, 3, 8]], [[0, 2], [2, 3], [3, 3]])expected [5, 8, 0]	<1 ms
  ✓	path_max(2, [[0, 1, 7]], [[1, 0]])expected [7]	<1 ms

  ✓ <1 ms✓ <1 ms✓ 1 ms✓ <1 ms✓ <1 ms⏱ 461 ms⏱ 573 ms
  Code

  import sys
  
  def path_max(n: int, edges: list[list[int]], queries: list[list[int]]) -> list[int]:
      # Build adjacency using arrays (CSR-like) for speed/memory
      m = n - 1
      deg = [0] * n
      for e in edges:
          deg[e[0]] += 1
          deg[e[1]] += 1
      start = [0] * (n + 1)
      for i in range(n):
          start[i + 1] = start[i] + deg[i]
      adj_v = [0] * (2 * m)
      adj_w = [0] * (2 * m)
      pos = start[:]
      for u, v, w in edges:
          p = pos[u]
          adj_v[p] = v
          adj_w[p] = w
          pos[u] = p + 1
          p = pos[v]
          adj_v[p] = u
          adj_w[p] = w
          pos[v] = p + 1
  
      # Iterative DFS: parent and max edge weight to parent
      parent = [-1] * n
      pw = [0] * n
      depth = [0] * n
      order = [0] * n  # preorder for binary lifting base
      stack = [0]
      parent[0] = 0
      idx = 0
      visited = [False] * n
      visited[0] = True
      while stack:
          u = stack.pop()
          order[idx] = u
          idx += 1
          pu = parent[u]
          du = depth[u]
          for i in range(start[u], start[u + 1]):
              v = adj_v[i]
              if not visited[v]:
                  visited[v] = True
                  parent[v] = u
                  pw[v] = adj_w[i]
                  depth[v] = du + 1
                  stack.append(v)
  
      # Binary lifting tables over preorder positions
      LOG = max(1, (n).bit_length())
      up = [parent]
      mx = [pw]
      for k in range(1, LOG):
          prev_up = up[k - 1]
          prev_mx = mx[k - 1]
          cur_up = [0] * n
          cur_mx = [0] * n
          for i in range(n):
              mid = prev_up[i]
              cur_up[i] = prev_up[mid]
              a = prev_mx[i]
              b = prev_mx[mid]
              cur_mx[i] = a if a > b else b
          up.append(cur_up)
          mx.append(cur_mx)
  
      res = [0] * len(queries)
      for qi in range(len(queries)):
          a = queries[qi][0]
          b = queries[qi][1]
          if a == b:
              res[qi] = 0
              continue
          ans = 0
          if depth[a] < depth[b]:
              a, b = b, a
          diff = depth[a] - depth[b]
          k = 0
          while diff:
              if diff & 1:
                  mw = mx[k][a]
                  if mw > ans:
                      ans = mw
                  a = up[k][a]
              diff >>= 1
              k += 1
          if a == b:
              res[qi] = ans
              continue
          for k in range(LOG - 1, -1, -1):
              ua = up[k][a]
              ub = up[k][b]
              if ua != ub:
                  mw = mx[k][a]
                  if mw > ans:
                      ans = mw
                  mw = mx[k][b]
                  if mw > ans:
                      ans = mw
                  a = ua
                  b = ub
          mw = mx[0][a]
          if mw > ans:
              ans = mw
          mw = mx[0][b]
          if mw > ans:
              ans = mw
          res[qi] = ans
      return res

• #1Z.ai GLM 5.3solved · 9 of 9 testsanswer in 5.7 s · runtime 1.1 s · $0.0013

  ✓	path_max(4, [[0, 1, 5], [1, 2, 3], [1, 3, 8]], [[0, 2], [2, 3], [3, 3]])expected [5, 8, 0]	<1 ms
  ✓	path_max(2, [[0, 1, 7]], [[1, 0]])expected [7]	2 ms

  ✓ 2 ms✓ 1 ms✓ <1 ms✓ <1 ms✓ <1 ms⏱ 635 ms⏱ 470 ms
  Code

  import sys
  
  def path_max(n, edges, queries):
      if n == 1 or not edges:
          return [0] * len(queries)
      adj = [[] for _ in range(n)]
      for u, v, w in edges:
          adj[u].append((v, w))
          adj[v].append((u, w))
  
      parent = [-1] * n
      pw = [0] * n
      order = []
      depth = [0] * n
      stack = [0]
      visited = [False] * n
      visited[0] = True
      while stack:
          u = stack.pop()
          order.append(u)
          for v, w in adj[u]:
              if not visited[v]:
                  visited[v] = True
                  parent[v] = u
                  pw[v] = w
                  depth[v] = depth[u] + 1
                  stack.append(v)
  
      size = [1] * n
      for u in reversed(order):
          for v, _ in adj[u]:
              if v != parent[u]:
                  size[u] += size[v]
  
      # heavy child: largest subtree
      heavy = [-1] * n
      for u in range(n):
          best = -1
          for v, _ in adj[u]:
              if v != parent[u] and (best == -1 or size[v] > size[best]):
                  best = v
          heavy[u] = best
  
      # iterative decomposition
      head = [0] * n
      pos = [0] * n
      cnt = 0
      # assign via iterative walk
      stack = [0]
      head[0] = 0
      while stack:
          u = stack.pop()
          pos[u] = cnt
          cnt += 1
          # push light children with their own chains first? use order:
          h = heavy[u]
          for v, _ in adj[u]:
              if v != parent[u]:
                  if v == h:
                      continue
                  head[v] = v
                  stack.append(v)
          if h != -1:
              head[h] = head[u]
              stack.append(h)
  
      # segment tree over pos, value pw[node]
      base = [0] * n
      for u in range(n):
          if u != 0:
              base[pos[u]] = pw[u]
      size2 = 1
      while size2 < n:
          size2 <<= 1
      NEG = 0
      seg = [NEG] * (2 * size2)
      seg[size2:size2 + n] = base
      for i in range(size2 - 1, 0, -1):
          seg[i] = seg[2 * i] if seg[2 * i] > seg[2 * i + 1] else seg[2 * i + 1]
  
      def query(l, r):
          res = 0
          l += size2
          r += size2 + 1
          while l < r:
              if l & 1:
                  if seg[l] > res:
                      res = seg[l]
                  l += 1
              if r & 1:
                  r -= 1
                  if seg[r] > res:
                      res = seg[r]
              l >>= 1
              r >>= 1
          return res
  
      def q(u, v):
          res = 0
          while head[u] != head[v]:
              if depth[head[u]] < depth[head[v]]:
                  u, v = v, u
              a = query(pos[head[u]], pos[u])
              if a > res:
                  res = a
              u = parent[head[u]]
          if u == v:
              return res
          if depth[u] > depth[v]:
              u, v = v, u
          a = query(pos[u] + 1, pos[v])
          if a > res:
              res = a
          return res
  
      return [q(a, b) for a, b in queries]

• #4Qwen3 Coder Nextsolved · 9 of 9 testsanswer in 12 s · runtime 1.4 s · $0.0006

  ✓	path_max(4, [[0, 1, 5], [1, 2, 3], [1, 3, 8]], [[0, 2], [2, 3], [3, 3]])expected [5, 8, 0]	1 ms
  ✓	path_max(2, [[0, 1, 7]], [[1, 0]])expected [7]	<1 ms

  ✓ <1 ms✓ 1 ms✓ 2 ms✓ <1 ms✓ 3 ms⏱ 553 ms⏱ 850 ms
  Code

  import sys
  from collections import defaultdict
  
  def path_max(n: int, edges: list[list[int]], queries: list[list[int]]) -> list[int]:
      if n == 1:
          return [0] * len(queries)
      
      # Build adjacency list
      graph = [[] for _ in range(n)]
      for u, v, w in edges:
          graph[u].append((v, w))
          graph[v].append((u, w))
      
      # Set up for binary lifting
      LOG = (n - 1).bit_length() if n > 1 else 1
      parent = [[-1] * n for _ in range(LOG)]
      max_edge = [[0] * n for _ in range(LOG)]
      depth = [0] * n
      
      # BFS/DFS to set up initial parent and depth information
      # Use iterative DFS to avoid recursion limit issues
      stack = [(0, -1, 0)]  # (node, parent, depth)
      while stack:
          node, par, dep = stack.pop()
          depth[node] = dep
          parent[0][node] = par
          for neighbor, weight in graph[node]:
              if neighbor != par:
                  max_edge[0][neighbor] = weight
                  stack.append((neighbor, node, dep + 1))
      
      # Preprocess binary lifting table
      for i in range(1, LOG):
          for node in range(n):
              if parent[i-1][node] != -1:
                  p = parent[i-1][node]
                  parent[i][node] = parent[i-1][p] if p != -1 else -1
                  max_edge[i][node] = max(max_edge[i-1][node], max_edge[i-1][p] if p != -1 else 0)
      
      def query_max(a, b):
          if a == b:
              return 0
          
          # Ensure a is deeper than b
          if depth[a] < depth[b]:
              a, b = b, a
          
          max_w = 0
          
          # Bring a up to the depth of b
          diff = depth[a] - depth[b]
          for i in range(LOG):
              if diff & (1 << i):
                  max_w = max(max_w, max_edge[i][a])
                  a = parent[i][a]
          
          if a == b:
              return max_w
          
          # Move both nodes up until their parents are the same
          for i in range(LOG - 1, -1, -1):
              if parent[i][a] != parent[i][b]:
                  max_w = max(max_w, max_edge[i][a], max_edge[i][b])
                  a = parent[i][a]
                  b = parent[i][b]
          
          # One more step to their LCA
          max_w = max(max_w, max_edge[0][a], max_edge[0][b])
          return max_w
      
      return [query_max(a, b) for a, b in queries]

Tokens cost: $0.0033 · How rounds are run and scored →