[LeetCode] Find the Pivot Integer

Problem

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.

주어진 n이 있을 때, (1, x)와 (x, n)의 합이 같은 경우가 되는 x를 구하는 문제이다.
만약 x가 존재하지 않는다면 -1을 반환한다.

Solution

먼저 prefix sum을 이용해서 직관적으로 코드를 작성했다.

1. 매번 범위 내의 합을 구하지 않기 위해 각 요소의 prefix sum을 구했다. 요소가 x라면 1부터 x의 합이 prefix_sum[x]에 저장된다.
2. n부터 1까지 순회하면서 현재 숫자가 x에 적합한지 확인한다. (1, x)는 prefix_sum[x]이고 (x, n)은 prefix_sum[n] - prefix_sum[x - 1]로 구할 수 있다. 두 값이 같다면 x를 반환한다.
3. 만약 모든 숫자를 순회했는데 x를 구하지 못한 경우에는 -1을 반환한다.
   코드는 아래와 같다.

class Solution:
    def pivotInteger(self, n: int) -> int:
        prefix_sum = [0] * (n + 1)
        for i in range(1, n + 1):
            prefix_sum[i] = prefix_sum[i - 1] + i
        for pivot in range(n, 0, -1):
            if prefix_sum[n] - prefix_sum[pivot - 1] == prefix_sum[pivot]:
                return pivot
        return -1

제출해보니 생각보다 시간복잡도가 많이 느렸다. 그래서 다른 사람의 discussion을 참고해서 수학적으로 접근한 코드를 새로 작성했다.

1. 수열의 합은 (1, n)일 때 $n(n + 1)/2$
2. 문제에서 필요한 합의 공식을 작성하면 아래와 같다.
   • (1, x) = $x(x + 1)/2$
   • (x, n) = (1, n) - (1, x - 1) = $n(n + 1)/2 - x(x - 1)/2$
3. (1, x) = (x, n)일 때의 x를 찾는 문제이므로 둘을 등식화하면 $x ^ 2 = n(n + 1)/2$라는 공식이 나온다.
4. x는 해당 값에 루트를 씌운 값이다.
5. x가 정수라면 x를 반환하고 정수가 아니라면 -1을 반환한다. 정수를 판별하는 방법은 x에서 x의 정수부와 차를 구했을 때 0이 아니라면 실수, 0이라면 정수이다.

class Solution:
    def pivotInteger(self, n: int) -> int:
        _sum = n * (n + 1) // 2
        x = math.sqrt(_sum)
        if x - math.ceil(x) == 0:
            return int(x)
        return -1

• 공식을 이용한 풀이는 시간복잡도가 O(1)으로 위의 코드의 시간복잡도 O(n)보다 빠르다.
• 디스커션 이름이 Math is everywhere인데 정말 맞는말이다 ㅎㅎ