Program to find sum of differences between max and min elements from randomly selected k balls from n balls in Python

Suppose we have n balls which are numbered by an array nums, whose size is n and nums[i] represents the number of ball i. Now we have another value k. In each turn we pick k balls from n different balls and find the difference of maximum and minimum values of k balls and store the difference in a table. Then put these k balls again into that pot and pick again until we have selected all possible selections. Finally find the sum of all differences from the table. If the answer is too large, then return result mod 10^9+7.

So, if the input is like n = 4 k = 3 nums = [5, 7, 9, 11], then the output will be 20 because combinations are −

• [5,7,9], difference 9-5 = 4
• [5,7,11], difference 11-5 = 6
• [5,9,11], difference 11-5 = 6
• [7,9,11], difference 11-7 = 4

so 4+6+6+4 = 20.

To solve this, we will follow these steps −

• m := 10^9 + 7
• inv := a new list with elements [0, 1]
• for i in range 2 to n, do
  • insert (m - floor of (m / i) * inv[m mod i] mod m) at the end of inv
• comb_count := 1
• res := 0
• for pick in range k - 1 to n - 1, do
  • res := res +(nums[pick] - nums[n - 1 - pick]) * comb_count mod m
  • res := res mod m
  • comb_count := comb_count *(pick + 1) mod m * inv[pick + 2 - k] mod m
• return res

Example

Let us see the following implementation to get better understanding −

def solve(n, k, nums):
   m = 10**9 + 7

   inv = [0, 1]
   for i in range(2, n + 1):
      inv.append(m - m // i * inv[m % i] % m)

   comb_count = 1
   res = 0
   for pick in range(k - 1, n):
      res += (nums[pick] - nums[n - 1 - pick]) * comb_count % m
      res %= m
      comb_count = comb_count * (pick + 1) % m * inv[pick + 2 - k] % m

   return res

n = 4
k = 3
nums = [5, 7, 9, 11]
print(solve(n, k, nums))

Input

4, 3, [5, 7, 9, 11]

Output

20