Beam Search over a Tree of Thoughts
Problem Statement
Chain-of-thought commits to one line of reasoning and lives with it. Tree of Thoughts generates several candidate thoughts at each step, scores them, keeps the best few and throws the rest away. The search itself is the algorithm - and it is a beam search with a pruning rule.
Background
You are given a pre-scored tree, so no model calls are needed. Each node has a score (the evaluator's judgement of that thought) and a list of children.
A path is a sequence of node ids from the root. Its value is the sum of the scores of its nodes excluding the root - the root is the problem statement, not a thought.
The search:
- Start with the beam holding the single root path, value 0.0.
- Repeat up to max_depth times:
- For each path in the beam, look at the last node's children. A path whose last node has no children is complete: move it to the finished set.
- Otherwise extend the path by each child. Every child examined counts towards expanded, whether or not it survives.
- Prune any child whose own score is < min_score; it never enters the beam.
- Collect all surviving extensions, sort by value descending (ties broken by the path itself, ascending, comparing ids element by element), and keep the first beam_width as the new beam.
- If nothing survives, the search stops.
- When the loop ends, any paths still in the beam join the finished set - they are depth-limited but still candidate answers.
- The answer is the finished path with the highest value, ties broken by the smallest path.
Your Task
Implement:
def tot_beam_search(nodes, root, beam_width, max_depth, min_score):
- nodes: dict of node id -> {"score": float, "children": list[str]} (children may be absent).
Return {"best_path": list[str], "best_value": float, "expanded": int} with best_value rounded to 4 decimal places. If nothing at all is finished, return [], 0.0 and the expansion count.
Input/Output Format
Compare values for sorting after rounding to 6 decimal places so that arithmetically-equal branches tie cleanly; report best_value rounded to 4.
Sample
nodes = {
"r": {"score": 0.0, "children": ["a", "b"]},
"a": {"score": 0.8, "children": []},
"b": {"score": 0.9, "children": []},
}
print(tot_beam_search(nodes, "r", 2, 3, 0.0))
# {'best_path': ['r', 'b'], 'best_value': 0.9, 'expanded': 2}
Example:
nodes = {'r': {'score': 0.0, 'children': ['a', 'b']}, 'a': {'score': 0.8, 'children': []}, 'b': {'score': 0.9, 'children': []}}
print(tot_beam_search(nodes, 'r', 2, 3, 0.0)){'best_path': ['r', 'b'], 'best_value': 0.9, 'expanded': 2}Both children of the root are examined, so expanded is 2. Neither is pruned (both scores clear min_score of 0.0) and both fit in a beam of width 2. On the next iteration each is childless, so both paths finish; ['r','b'] has value 0.9 against 0.8 and wins.
Constraints:
1 <= len(nodes) <= 500; the tree is acyclic and every child id exists innodes- The root's own score never contributes to any path value
expandedcounts every child examined, including pruned ones- Pruning compares the child's own score against
min_score, not the path value - Sorting compares values rounded to 6 dp, then the path list ascending
best_valueis rounded to 4 decimal places - never print a raw float- Paths still in the beam when
max_depthis exhausted are valid answers
1. Background Knowledge
Beam Search is a heuristic search algorithm that explores a graph by expanding the most promising nodes in a limited set, known as the beam. Unlike breadth-first search, which keeps all nodes at the current depth, beam search maintains only the top k nodes (where k is the beam width). This makes it memory-efficient and suitable for problems where the search space is too large for exhaustive exploration, such as natural language generation or planning tasks.
In the context of Tree of Thoughts (ToT), the search space is a tree where each node represents a "thought" or a step in reasoning. Each node has an associated score reflecting its quality or likelihood of leading to a correct solution. The goal is to find the path from the root to a leaf (or a depth-limited node) that maximizes the cumulative score. The pruning step, controlled by min_score, allows the algorithm to discard low-quality branches early, focusing computational resources on the most promising paths.
The value of a path is defined as the sum of the scores of all nodes in the path, excluding the root. This distinction is crucial because the root represents the initial problem statement, not a generated thought. The search proceeds iteratively, expanding paths, pruning low-scoring children, and selecting the best candidates for the next iteration. The process continues until the maximum depth is reached or no further expansions are possible.
2. Algorithm Approach
The core algorithm is a modified beam search with explicit pruning and depth control. The approach involves maintaining a beam of active paths, each represented as a list of node IDs. At each step, the algorithm expands the last node of each path in the beam, evaluates the children, prunes those below min_score, and selects the top beam_width paths based on their cumulative value.
The algorithm uses a loop that runs up to max_depth times. In each iteration, it processes all paths in the current beam. For each path, it checks if the last node has children. If not, the path is considered complete and moved to a finished set. If it has children, each child is examined, and its score is checked against min_score. Surviving children are added to a new beam, and their cumulative values are updated. After processing all paths, the new beam is sorted by value (descending) and path ID (ascending for ties), and truncated to beam_width.
The search terminates when the loop completes or the beam becomes empty. The final result is the path with the highest value among all finished paths, including those that reached max_depth. If no paths are finished, the algorithm returns an empty path and a value of 0.0.
3. Step-by-Step Strategy
- Initialize: Create a beam list containing a single path [root] with value 0.0. Initialize expanded to 0 and finished_paths to an empty list.
- Loop: Iterate up to max_depth times.
- Expand: For each path in the beam:
- Get the last node ID.
- If the node has no children, add the path to finished_paths and continue.
- Otherwise, iterate through each child. Increment expanded for each child examined.
- If the child's score is ≥ min_score, create a new path by appending the child ID. Calculate the new value as the parent path's value plus the child's score. Add this new path and value to a temporary list next_beam.
- Prune and Select: Sort next_beam by value (descending) and path ID (ascending). Keep only the first beam_width entries. Update beam to this sorted list.
- Check Termination: If beam is empty, break the loop.
- Finalize: After the loop, add any remaining paths in beam to finished_paths.
- Select Best: If finished_paths is empty, return [], 0.0, and expanded. Otherwise, find the path with the highest value (ties broken by smallest path ID). Round the best value to 4 decimal places and return the result.
4. Common Pitfalls
- Root Score Exclusion: Remember that the root's score is not included in the path value. The initial value is 0.0, and only children's scores are added.
- Expansion Count: The expanded counter must increment for every child examined, regardless of whether it survives pruning. This includes children that are pruned due to min_score.
- Tie-Breaking: When sorting paths, ties in value must be broken by comparing path IDs element by element in ascending order. This ensures deterministic behavior.
- Depth Limit: Paths that reach max_depth are not necessarily leaves. They should be added to finished_paths at the end of the loop, not during expansion.
- Empty Beam: If the beam becomes empty during the loop, the search should stop immediately. Do not continue iterating if there are no paths to expand.
- Rounding: Values should be compared after rounding to 6 decimal places to handle floating-point precision issues, but the final best_value should be rounded to 4 decimal places.
5. Time & Space Complexity
The time complexity is O(D⋅B⋅C⋅log(B)), where D is max_depth, B is beam_width, and C is the average number of children per node. In each of the D iterations, we process up to B paths, each with up to C children. Sorting the next_beam takes O(B⋅C⋅log(B⋅C)), which simplifies to O(B⋅C⋅log(B)) since C is typically small. The expanded count is O(D⋅B⋅C).
The space complexity is O(B⋅D) to store the beam and finished paths. Each path has length up to D, and we store up to B paths in the beam and a similar number in finished_paths. The nodes dictionary is given and not counted in the auxiliary space.