How many balls do you need?

#Beta
#https://www.codewars.com/kata/57f945b77a28db0d750001f2/python

def count_balls(n):
    balls = 0
    s = 1
    for i in range(1, n + 1):
        balls += s
        s += i + 1
    return balls