1. 문제 소개

3354. Make Array Elements Equal to Zero

2. 나의 풀이법

Easy난이도라 간단한 풀이였다.
값이 0인 인덱스를 고르고 그 인덱스를 기준으로 왼쪽합과 오른쪽합이 같은지, 차이가 1인지, 마지막으로 다른 지에 대해서 알아보면 되었다.
하지만 첫 풀이는 반복적으로 sum()을 호출하므로 시간복잡도는 $O(n^2)$가 되어 비효율적이다.
prefix와 suffix로 변경하여 시간을 조금 더 단축했다.
이 경우에는 시간복잡도가 $O(n)$으로 줄어든다.

class Solution:
    def countValidSelections(self, nums: List[int]) -> int:
        n = len(nums)
        ans = 0
        prefix, total = 0, sum(nums)

        for i in range(n):
            left, right = prefix, total - prefix - nums[i]
            if nums[i] == 0:
                if left == right:
                    ans += 2
                elif abs(left - right) == 1:
                    ans += 1
            prefix += nums[i]
        return ans

3. 다른 풀이법

다른 풀이의 경우 문제의 알고리즘을 그대로 구현했다.
시간복잡도는 sum()을 반복해서 호출하는 코드와 똑같은 $O(n^2)$이지만 실제로 푸는 시간은 차이가 있다.

class Solution:
    def countValidSelections(self, nums: List[int]) -> int:
        def simulate(i, d):
            arr = nums[:]
            while 0 <= i < len(arr):
                if arr[i] == 0:
                    i += d
                else:
                    arr[i] -= 1
                    d *= -1
                    i += d
            return all(x == 0 for x in arr)

        res = 0
        for i in range(len(nums)):
            if nums[i] == 0:
                if simulate(i, 1): res += 1
                if simulate(i, -1): res += 1
        return res

4. 결론

같은 결과를 나타내는 코드더라도 작동원리를 파악하고 단순화하는 것이 중요함을 깨달았다.