1712 Ways to Split Array Into Three Subarrays

숫자로 구성된 배열을 각각 mid의 합은 left의 합보다 같거나 크게, right의 합은 mid의 합보다 같거나 크게 3조각으로 나눌 수 있는 경우의 수를 구하는 문제이다. 이 문제는 정렬된 배열에서 특정한 값을 찾는 전형적인 이진 탐색 문제로,

누적합이 저장된 배열을 만들고, 그 배열에서 mid가 될 수 있는 최소 인덱스와 최대 인덱스를 찾으면 된다.

그리고 이진 탐색 말고도, 투 포인터로 풀린다. 나도 처음에 이진 탐색 말고도, 투 포인터로 하려고 했는데 예외 처리에서

애먹어서 그냥 이진 탐색으로 풀었다. 코드는 Discuss 뒤지면 나온다.

 

C++

class Solution {
public:
    int waysToSplit(vector<int>& nums) {
        vector<int> sum(nums.size());
        vector<int>::iterator lb, ub;
        
        sum[0] = nums[0];
        for(size_t i=1;i<nums.size();++i)
            sum[i] = sum[i - 1] + nums[i];
        
        int cnt{};
        for(size_t l = 0; l < sum.size() - 2; ++l){
            lb = lower_bound(sum.begin() + l + 1, sum.end(), sum[l] * 2);
            ub = upper_bound(sum.begin() + l + 1, sum.end() - 1, (sum.back() - sum[l]) / 2 + sum[l]);
            cnt = (cnt + max(int(ub - lb), 0)) % 1000000007;
        }
        return cnt;
    }
};

Python

class Solution:
    def waysToSplit(self, nums: List[int]) -> int:
        sums = [nums[0]]
        for n in nums[1:]:
            sums.append(n + sums[-1])
        
        cnt = 0
        for l in range(len(sums) - 2):
            lb = bisect_left(sums, sums[l] * 2)
            ub = bisect_right(sums, (sums[-1] + sums[l])/2)
            cnt = (cnt + max(0, min(ub, len(sums)-1) - max(lb, l+1))) % 1000000007
        return cnt​