HackerRank Combinatorics – Choose and Calculate

Challenge: Choose and Calculate

Subdomeniu: Combinatorics (combinatorics)

Scor cont: 40.0 / 40

Submission status: Accepted

Submission score: 1.0

Submission ID: 464728853

Limbaj: python3

Link challenge: https://www.hackerrank.com/challenges/choose-and-calculate/problem

Cerinta

Animesh and Mohit are playing a game. They have $N$ balls in front of them. All the balls are numbered, not necessarily in any order. Animesh picks a set of $K$ balls from the lot each time, checks this set, puts it back into the lot, and repeats this process until all possible selections have been made. At each draw, Mohit picks two balls from these $K$ balls -- the one with the maximum number and the one with the minimum number, and notes their difference on a paper. At the end of the game, Animesh has to tell the $Sum$ of all the numbers on Mohit's paper. As value of this number can be very huge, you have to tell the value _mod_ $10^9+7$.  

**Input Format**  
The first line contains two integers $N$ and $K$.  
The next line will contain a list having $N$ numbers.  

**Output Format**  
Print the value of $Sum$  _mod_ $(10^9+7)$.  

**Constraints**  
$1 \le N \le 10^5$  
$1 \le K \le N$  
$0 \le \text{ numbers_on_ball } \le 10^9$  

**Sample Input**  

	4 2
	10 20 30 40

    
**Sample Output**  

	100

**Explanation**

There are 6 possible selections.  
1. 10 20  
2. 20 30  
3. 30 40  
4. 10 30  
5. 20 40  
6. 10 40  
Summation = 10+10+10+20+20+30 = 100

Cod sursa

#!/bin/python3

import math
import os
import random
import re
import sys

def solve(balls, k):
    MOD = 10**9+7
    result = 0
    n = len(balls)
    balls.sort()
    
    memo = {}
    memo[k-1] = math.comb(k-1, k-1)
    for i in range(k, n+1):
        memo[i] = memo[i-1] * i // (i-(k-1))
    for i in range(n-k+1):
        result -= balls[i] * memo[n-1-i]
    for i in range(k-1, n):
        result += balls[i] * memo[i]
    return result % MOD
if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    first_multiple_input = input().rstrip().split()

    n = int(first_multiple_input[0])

    k = int(first_multiple_input[1])

    balls = list(map(int, input().rstrip().split()))

    result = solve(balls, k)

    fptr.write(str(result) + '\n')

    fptr.close()