fork download
copy 
  1. # your code goes here
  2. def palindromePaths(tree_nodes, tree_from, tree_to, arr, queries):
  3.     # 1. Build the adjacency list for the undirected tree
  4.     adj = [[] for _ in range(tree_nodes)]
  5.     for u, v in zip(tree_from, tree_to):
  6.         adj[u].append(v)
  7.         adj[v].append(u)
  8.  
  9.     # 2. Array to store the precalculated answer for every node
  10.     node_answers = [0] * tree_nodes
  11.  
  12.     # 3. Hash map to store the frequency of masks of the ancestors
  13.     freq = {}
  14.  
  15.     # 4. Stack for Iterative DFS to avoid RecursionError on deep trees.
  16.     # Elements: (current_node, parent_node, path_mask_to_parent, state)
  17.     # State 0: Entering the node (pre-order)
  18.     # State 1: Exiting the node (post-order, for backtracking)
  19.     stack = [(0, -1, 0, 0)]
  20.  
  21.     while stack:
  22.         u, p, current_mask, state = stack.pop()
  23.  
  24.         if state == 0:
  25.             # Toggle the bit corresponding to the character at node `u`
  26.             char_idx = ord(arr[u]) - ord('a')
  27.             P_u = current_mask ^ (1 << char_idx)
  28.  
  29.             # Add the parent's path mask to our active ancestor frequencies
  30.             freq[current_mask] = freq.get(current_mask, 0) + 1
  31.  
  32.             # Count valid ancestors v where the path from u to v is a palindrome
  33.             # Condition: P_u ^ P_parent(v) has at most 1 bit set (0 or power of 2)
  34.  
  35.             # Case 1: 0 bits set (perfectly even frequencies)
  36.             count = freq.get(P_u, 0)
  37.  
  38.             # Case 2: exactly 1 bit set (one odd character frequency)
  39.             for i in range(26):
  40.                 count += freq.get(P_u ^ (1 << i), 0)
  41.  
  42.             # Save the answer for the current node
  43.             node_answers[u] = count
  44.  
  45.             # Push the backtrack state for the current node
  46.             stack.append((u, p, current_mask, 1))
  47.  
  48.             # Push children to the stack
  49.             for v in adj[u]:
  50.                 if v != p:
  51.                     stack.append((v, u, P_u, 0))
  52.         else:
  53.             # Backtrack: Remove the current ancestor mask as we exit the branch
  54.             freq[current_mask] -= 1
  55.             if freq[current_mask] == 0:
  56.                 del freq[current_mask]
  57.  
  58.     # 5. Return answers for the required queries in O(1) time each
  59.     return [node_answers[q] for q in queries]
  60.  
  61.  
  62. print(palindromePaths(7, [0,1,2,2,4,4,7], [1,2,3,4,5,6], ['a', 'b', 'c', 'a', 'c','b', 'c'], [6,5,3]))
Success #stdin #stdout 0.09s 14092KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
[3, 4, 1]
https://ideone.com/NlB18E
language:
Python 3 (python  3.12)
created:
3 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!