Problem

Given a non-negative integer number n, we need to find the pivot number. The pivot number p is a number such as the sum of all numbers [0..p] equals the sum of all numbers [p..n].

Examples:

Solution 1

fn find_pivot_number(n: u32) -> Option<u32> {    
    let sum: u32 = n * (n + 1) / 2;

    let mut left_sum = 0;
    (0..=n).find_map(|e| {
        let right_sum = sum - left_sum;
        left_sum += e;
        if right_sum == left_sum {
            Some(e)
        } else {
            None
        }
    })
}

Solution 2

fn find_pivot_number(n: u32) -> Option<u32> {    
    let sum = n * (n + 1) / 2;
    let p = (sum as f32).sqrt() as u32;
    (p * p == sum).then(|| p)
}

Explanation

As per the Arithmetic progression formula, we know that the sum of the first n integer numbers could be calculated as: n * (n + 1) / 2. Thus, we can calculate sum without iterating over the numbers both in Solution 1 and Solution 2.

Then in Solution 1, we could iterate over numbers [0..n] maintaining the current left_sum. And since we know the overall sum of the numbers, we could easily figure out the right_sum = sum - left_sum. If left_sum and right_sum match, then we have found the pivot. Solution 1 is OK but it has O(n) time complexity, which could be improved...

Surprisingly (or not), we could solve the puzzle more efficiently and concisely just by using the sqrt function in Solution 2. It is easy to prove that sqrt would give us exactly what we need. We know that the sum of all numbers sum[0..n] = n * (n + 1) / 2. Then sum[0..p] = p * (p + 1) / 2 and sum[p..n] = sum[0..n] - (p * (p + 1) / 2) + p. We need to find a number such as sum[0..p] = sum[p..n], i.e. p * (p + 1) / 2 = sum[0..n] - (p * (p + 1) / 2) + p. If we simplify this equation, that would give us sum[0..n] = p * p. Which means p = sqrt(sum[0..n]).

Complexity

Solution 1: The time complexity of the solution is O(n) (where n is the input number).

Solution 2: The time complexity of the solution is O(1).

The auxiliary space complexity is O(1) for both solutions.

---

Rust Playground with this code