tickets.py

gistfile1.py

# Enter your code here. Read input from STDIN. Print output to STDOUT
n, m = raw_input().split(" ")
n = int(n)
m = int(m)
tickets = map(lambda x: int(x), raw_input().split(" "))

# n, m = 2, 4
# tickets = [5, 2]

sorted_tickets = sorted(tickets, reverse=True)
tickets_available = m

def sum_progression(a1, an):
    return (a1+an+1)*(an-a1)/2 # Is the sum of a aritmetic progression http://en.wikipedia.org/wiki/Arithmetic_progression
    # return sum(range(a1+1, n+1)) ~ Very bad performant

i = 0
ac = 0
max_ti = sorted_tickets[i]

while tickets_available>0:
    if i == len(sorted_tickets):
        # We are in the last element
        diff = max_ti
    else:
        current = sorted_tickets[i]
        if current == max_ti:
            # We want to group all the same rows
            i += 1
            continue
        diff = max_ti-current
    
    retire_tickets = min(diff*i, tickets_available)
    if retire_tickets == 0:
        # we cannot take more tickets
        break
    rows = min(retire_tickets // i, diff)
    cols = retire_tickets % i
    ac += sum_progression(max_ti-rows, max_ti)*i + cols*(max_ti-rows)
    max_ti -= rows
    tickets_available -= rows*i+cols

print ac