Teaching Kids Programming – Algorithms to Find the Pivot Integer of the First N Natural Numbers

Teaching Kids Programming: Videos on Data Structures and Algorithms

Given a positive integer n, find the pivot integer x such that: The sum of all elements between 1 and x inclusively equals the sum of all elements between x and n inclusively. Return the pivot integer x. If no such integer exists, return -1. It is guaranteed that there will be at most one pivot index for the given input.

Example 1:
Input: n = 8
Output: 6
Explanation: 6 is the pivot integer since: 1 + 2 + 3 + 4 + 5 + 6 = 6 + 7 + 8 = 21.

Example 2:
Input: n = 1
Output: 1
Explanation: 1 is the pivot integer since: 1 = 1.

Example 3:
Input: n = 4
Output: -1
Explanation: It can be proved that no such integer exist.

Constraints:
1 <= n <= 1000
Can you use brute force to check every number from 1 to n if any of them is the pivot integer?
If you know the sum of [1: pivot], how can you efficiently calculate the sum of the other parts?

Find the Pivot Integer of the First N Natural Numbers (Bruteforce, Binary Search, Math, Prefix Sum)

We can solve this puzzle via four different algorithms:

Bruteforce Algorithm

We can certainly bruteforce the possible pivot integer X – there are only N choices. We try them until we find one.

1
2
3
4
5
6
7
8
9

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        for i in range(1, n + 1):
            a = sum(range(1, i + 1))
            b = sum(range(i, n + 1))
            if a == b:
                return i
        return -1

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        for i in range(1, n + 1):
            a = sum(range(1, i + 1))
            b = sum(range(i, n + 1))
            if a == b:
                return i
        return -1

Bruteforce Algorithm via Next Function

This can be also rewritten using the next function in one line:

1
2
3
4

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        return next((i for i in range(1, n + 1) if sum(range(1, i + 1)) == sum(range(i, n + 1))), -1 )

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        return next((i for i in range(1, n + 1) if sum(range(1, i + 1)) == sum(range(i, n + 1))), -1 )

The time complexity is O(N^2). It is quadratic because the sum inside the loop takes O(N) time.

Prefix Sum Algorithm (accumulated sum)

We can bring down the overall time complexity to O(N) by making the calculation of left and right part using the accumulate or prefix sum:

1
2
3
4
5
6
7
8
9
10
11

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:        
        a = 0
        total = sum(range(1, n + 1))
        for i in range(1, n + 1):
            a += i
            b = total - a + i
            if a == b:
                return i
        return -1

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:        
        a = 0
        total = sum(range(1, n + 1))
        for i in range(1, n + 1):
            a += i
            b = total - a + i
            if a == b:
                return i
        return -1

The time complexity is O(N), and the space complexity is O(1) constant. Here is O(N) space to store the prefix sums using the accumulate function:

1
2
3
4
5
6
7
8
9
10
11

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        total = sum(range(1, n + 1))
        prefix = list(accumulate(range(1, n + 1)))
        for i in range(1, n + 1):
            a = prefix[i - 1]
            b = total - prefix[i - 1] + i
            if a == b:
                return i
        return -1

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        total = sum(range(1, n + 1))
        prefix = list(accumulate(range(1, n + 1)))
        for i in range(1, n + 1):
            a = prefix[i - 1]
            b = total - prefix[i - 1] + i
            if a == b:
                return i
        return -1

Math Algorithm via Equation of Arithmetic Progression

Since it is consecutive numbers, we can compute the sum using the math equation for sum of any arithmetic progression numbers:

1
2
3
4
5
6
7
8
9

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        for i in range(1, n + 1):
            a = (1 + i) * i // 2
            b = (i + n) * (n - i + 1) // 2
            if a == b:
                return i
        return -1

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        for i in range(1, n + 1):
            a = (1 + i) * i // 2
            b = (i + n) * (n - i + 1) // 2
            if a == b:
                return i
        return -1

The sum of a, a + 1, a + 2, … b can be computed as:

.

This can be proved using Math Induction.

Binary Search Algorithm

This has O(N) time and O(1) space. We can binary search the pivot integer so that the time complexity is

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        L = 1
        R = n
        while L <= R:
            m = L + R >> 1
            a = (1 + m) * m // 2
            b = (m + n) * (n - m + 1) // 2
            if a == b:
                return m
            if a < b:
                L = m + 1
            else:
                R = m - 1
        return -1

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        L = 1
        R = n
        while L <= R:
            m = L + R >> 1
            a = (1 + m) * m // 2
            b = (m + n) * (n - m + 1) // 2
            if a == b:
                return m
            if a < b:
                L = m + 1
            else:
                R = m - 1
        return -1

If the sum of Left is smaller than sum of Right, we need to move the pivot index to the right, otherwise move the pivot index to the left.

Math Proof to Solve the Equation

Let’s assum the pivot integer is x. The left sum is , the right sum is the total sum minus the left sum plus x which is

So, we need to check if the square root of the sum between 1 to N inclusive is a square root. The space complexity is O(1) constant, and the time complexity is O(1) amortized since the modern CPUs know how to compute the square roots efficiently and has inbuilt specific instructions to handle the SQRT calculations.

Obviously, this is the most efficient solution among all.

1
2
3
4
5
6
7

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        x = sqrt((n+1)*n//2)
        if x.is_integer():
            return int(x)
        return -1

class Solution:
    @cache
    def pivotInteger(self, n: int) -> int:
        x = sqrt((n+1)*n//2)
        if x.is_integer():
            return int(x)
        return -1

–EOF (The Ultimate Computing & Technology Blog) —