We use cookies to ensure you have the best browsing experience on our website. Please read our cookie policy for more information about how we use cookies.
  1. Prepare
  2. Mathematics
  3. Combinatorics
  4. Choose and Calculate
  5. Discussions

Choose and Calculate

  • kohnicholas806
     + 0 comments

    Hi everyone, I get TLE for the last 4 test cases... please advise me on how to optimize the code. Thank you!

    >

    mod = 10**9 + 7 def solve(balls, k): if k == 1: return 0

    balls.sort()
ans = 0
n = len(balls)

for i, ball in enumerate(balls):
    if i >= k-1:
        ans += ball * (comb(i, k-1))

    if i < n-k+1:
        ans -= ball * comb(n-i-1, k-1)

return ans % mod