Leetcode # 2867. Count Valid Paths in a Tree

Problem

https://leetcode.com/contest/weekly-contest-364/problems/count-valid-paths-in-a-tree/

• Notice! The relation of a and b in path(a, b) is not only relation between ancestor and descendant

Testcases

相關例題

Solution: DFS + Sliding Window
[Time Limit Exceeded  717 / 922 testcases passed]

Time Complexity: O()
Space Complexity: O()
(The input and output generally do not count towards the space complexity.)

class Solution:
  def countPaths(self, n: int, edges: List[List[int]]) -> int:
    nums = [True] * (n + 1)
    primes = set()
    for i in range(2, n + 1):
      if nums[i]:
        primes.add(i)
        for j in range(2 * i, n + 1, i):
          nums[j] = False

    _edges = defaultdict(list)
    for u, v in edges:
      _edges[u].append(v)
      _edges[v].append(u)
    
    # DFS
    paths = {}
    for u in range(1, n + 1):
      stack, visited = [[[u], [0] if u in primes else []]], {u}
      while stack:
        curr, prime_i = stack.pop()
        is_leaf = True
        for child in _edges[curr[-1]]:
          if child in visited: continue
          is_leaf = False
          stack.append([curr + [child], 
                        prime_i + ([len(curr)] if child in primes else [])])
        visited.add(curr[-1])
        if is_leaf:
          pair = (min(curr[0], curr[-1]), max(curr[0], curr[-1]))
          if pair not in paths:
            paths[pair] = [curr, [-1] + prime_i + [len(curr)]]
    
    # sliding window
    pairs = set()
    for item in paths.items():
      path, prime_i = item[1]
      for p in range(1, len(prime_i) - 1):
        for l in range(prime_i[p - 1] + 1, prime_i[p] + 1):
          for r in range(prime_i[p], prime_i[p + 1]):
            arr = path[l:r + 1]
            if len(arr) < 2: continue
            pair = (min(arr[0], arr[-1]), max(arr[0], arr[-1]))
            if pair in pairs: continue
            pairs.add(pair)
            
    return len(pairs)

Last Updated on 2023/10/04 by A1go