[#6] Sum square difference

The sum of the squares of the first ten natural numbers is,

1² + 2² + ... + 10² = 385

The square of the sum of the first ten natural numbers is,

(1 + 2 + ... + 10)² = 55² = 3025

Hence the difference between the sum of the squares of the first ten natural numbers and the square of the sum is 3025 − 385 = 2640.
Find the difference between the sum of the squares of the first one hundred natural numbers and the square of the sum.

def problem6(n):
        '''
        The sum of the squares of the first ten natural numbers is,
        1^2 + 2^2 + ... + 10^2 = 385
        The square of the sum of the first ten natural numbers is,
        (1 + 2 + ... + 10^)2 = 55^2 = 3025
        Hence the difference between the sum of the squares of the first ten natural numbers
        and the square of the sum is 3025 − 385 = 2640.
        Find the difference between the sum of the squares of the first n natural numbers
        and the square of the sum.
        '''
        sum_of_square = n * (n + 1) * (2 * n + 1) // 6 #1^2 + 2^2 + ... + n^2 = n * (n + 1) * (2 * n + 1) // 6
        square_of_sum = (n * (n + 1) // 2) ** 2 #1 + 2 + ... + n = n * (n + 1) // 2
        return square_of_sum - sum_of_square