Skip to content

Commit 4e817a8

Browse files
ivnvxdcclauss
andauthored
Add boundary checking to optimal bst (#11771)
* Add boundary checking to optimal bst * Add additional checks to prevent IndexError --------- Co-authored-by: Christian Clauss <cclauss@me.com>
1 parent d182474 commit 4e817a8

1 file changed

Lines changed: 12 additions & 6 deletions

File tree

dynamic_programming/optimal_binary_search_tree.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,7 @@ def find_optimal_binary_search_tree(nodes) -> None:
102102
# This 2D array stores the overall tree cost (which's as minimized as possible);
103103
# for a single key, cost is equal to frequency of the key.
104104
dp = [[freqs[i] if i == j else 0 for j in range(n)] for i in range(n)]
105-
# sum[i][j] stores the sum of key frequencies between i and j inclusive in nodes
105+
# total[i][j] stores the sum of key frequencies between i and j inclusive in nodes
106106
# array
107107
total = [[freqs[i] if i == j else 0 for j in range(n)] for i in range(n)]
108108
# stores tree roots that will be used later for constructing binary search tree
@@ -115,11 +115,17 @@ def find_optimal_binary_search_tree(nodes) -> None:
115115
dp[i][j] = sys.maxsize # set the value to "infinity"
116116
total[i][j] = total[i][j - 1] + freqs[j]
117117

118-
# Apply Knuth's optimization
119-
# Loop without optimization: for r in range(i, j + 1):
120-
for r in range(root[i][j - 1], root[i + 1][j] + 1): # r is a temporal root
121-
left = dp[i][r - 1] if r != i else 0 # optimal cost for left subtree
122-
right = dp[r + 1][j] if r != j else 0 # optimal cost for right subtree
118+
# Apply Knuth's optimization with safe boundary handling
119+
r_start = root[i][j - 1] if j > i else i
120+
r_end = root[i + 1][j] if i < j else j
121+
122+
# Ensure r_start and r_end are within valid bounds
123+
r_start = max(i, min(r_start, j))
124+
r_end = min(j, max(r_end, i))
125+
126+
for r in range(r_start, r_end + 1):
127+
left = dp[i][r - 1] if r > i else 0 # optimal cost for left subtree
128+
right = dp[r + 1][j] if r < j else 0 # optimal cost for right subtree
123129
cost = left + total[i][j] + right
124130

125131
if dp[i][j] > cost:

0 commit comments

Comments
 (0)