LeetCode/rs/find-the-pivot-integer.rs

struct Solution {}
impl Solution {
    pub fn pivot_integer(n: i32) -> i32 {
        // S_i = S_n - S_{i - 1}
        // i * (1 + i) / 2 = n * (1 + n) / 2 - i * (i - 1) / 2
        // i * (1 + i) = n * (1 + n) - i * (i - 1)
        // 2i² = n(1 + n)
        // i = sqrt(n * (1 + n) / 2)

        let expected_ii = n * (1 + n) / 2;

        let mut i = 1;
        while i * i < expected_ii {
            i += 1;
        }

        if i * i == expected_ii {
            return i;
        }

        -1
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn example_1() {
        assert_eq!(Solution::pivot_integer(8), 6);
    }

    #[test]
    fn example_2() {
        assert_eq!(Solution::pivot_integer(1), 1);
    }

    #[test]
    fn example_3() {
        assert_eq!(Solution::pivot_integer(4), -1);
    }
}