Follow

Follow

# Rust Interview Puzzles: Find the pivot number

go4fun.fun
·Jan 30, 2023·

### 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:

 Input `n` Output Explanation `1` `1` `4` `None` No pivot `5` `None` No pivot `8` `6` `sum[0..6] = sum[6..8] = 21` `49` `35` `sum[0..35] = sum[35..49] = 630`

### 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