Skip to content

Commit 55cf4d6

Browse files
committed
Fix intermediate integer overflow in math.perm and math.comb
1 parent 6b00cba commit 55cf4d6

1 file changed

Lines changed: 26 additions & 4 deletions

File tree

src/runtime/math.py

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -225,11 +225,27 @@ def comb(n: i32, k: i32) -> i32:
225225
Computes the result of `nCk`, i.e, the number of ways to choose `k`
226226
items from `n` items without repetition and without order.
227227
"""
228-
229-
if n < k or n < 0:
228+
if k < 0 or k > n:
230229
return 0
231-
return i32(factorial(n)//(factorial(k)*factorial(n-k)))
230+
if k == 0 or k == n:
231+
return 1
232+
233+
# Creating a local variable because assignment to input function parameter 'k' is not allowed
234+
k_final: i32 = k
235+
236+
# Optimization: nCk is the same as nC(n-k), (done so that the number of iterations in loop can be reduced).
237+
if k_final > n // 2:
238+
k_final = n - k_final
239+
240+
result: i32 = 1
241+
i: i32
242+
243+
for i in range(k_final):
244+
# Calculate: result * (n - i) / (i + 1)
245+
result = result * (n - i)
246+
result = result // (i + 1)
232247

248+
return result
233249

234250
def perm(n: i32, k: i32) -> i32:
235251
"""
@@ -239,7 +255,13 @@ def perm(n: i32, k: i32) -> i32:
239255

240256
if n < k or n < 0:
241257
return 0
242-
return i32(factorial(n)//factorial(n-k))
258+
if k == 0:
259+
return 1
260+
result: i32 = 1
261+
i: i32
262+
for i in range(n, n-k, -1):
263+
result = result * i
264+
return result
243265

244266

245267
def isqrt(n: i32) -> i32:

0 commit comments

Comments
 (0)