[LeetCode 1712] Ways to Split Array Into Three Subarrays ( C++/Python)

문제 요약 

array를 3분할 (left, mid, right) left의 sum <= mid의 sum <= right의 sum 이 되게 하는 경우의 수를 구하라.

 

아이디어

sum값을 미리 만들어둔다

[1,1,1] --> [1,2,3] 

[1,2,3,4,5] --> [1,3,6,10,15] 이런식으로.

처음에 그냥 탐색으로 풀었다가 Timeout났었고, binary search로 접근하다가 조건이 뭔지 몰라서 한참 헤맸다.

결국 유튜브에 설명한번 보고 binary search로 해결했다.

python 치트키인 bisect로 풀었길레, C++로 재구현해서 C++/python모두 제출했다.

 

C++(시간초과 코드)

class Solution {
public:
    int dp[100001];
    int waysToSplit(vector<int>& nums) {
        int left,right;
        int last = nums.size()-1;
        left = 0;
        //dp계산.
        dp[0] = nums[0];
        for(int i = 1 ; i<=last; i++)
        {
            dp[i] = dp[i-1]+nums[i];
        }
        
        int count = 0;
        for(left = 0 ; left<= last-2; left++)
        {
            for(right = left+1 ; right<= last-1; right++)
            {
                int block1 = dp[left];
                int block2 = dp[right]-block1;
                int block3 = dp[last]-dp[right];
                if(block1 <= block2 && block2 <= block3) count++;
            }
        }
        return count;
    }
};

75개정도 통과하고 timeout난다.

 

임의의 array가 있을 때 left까지가 첫번쨰 블록이고, right-left가 두번째 블럭, 마지막 까지의 합(last) - right가 세번쨰 블럭이 되어서 조건 만족할때마다 count++하면서 풀었었다.

 

C++ 풀이(이진탐색)

#include <algorithm>
class Solution {
public:
    int dp[100001];
    int waysToSplit(vector<int>& nums) {
        dp[0] = nums[0];
        for(int i = 1 ; i< nums.size(); i++)
        {
            dp[i] = dp[i-1]+nums[i];
        }
        
        int cnt = 0;
        for(int i = 0 ; i< nums.size(); i++)
        {
            int lower,upper;
            //lower찾고 dp[i]*2 
            int left,right;
            left = i+1;
            right = nums.size()-1;
            while(left < right)
            {
                int mid = left+(right-left)/2;
                
                //>= 부등호 주의.
                if(dp[mid] >= dp[i]*2) 
                {
                    right = mid;
                }else{
                    left = mid+1;
                }
            }
            lower = left;
            //-----------------------------------------
            
            //upper찾고 i와 마지막점의 합 나누기 2.
            left = i+1;
            right = nums.size()-1;
            while(left < right)
            {
                int mid = left+(right-left)/2;
                
                //>= 부등호 주의.
                if((dp[nums.size()-1]+dp[i])/2 < dp[mid])
                {
                    right = mid;
                }else{
                    left = mid+1;
                }
            }
            upper = left;
            //---------------------------------------------
            cnt += max(0, min(upper, (int)nums.size()-1) - max(i+1, lower) );
            cnt %= 1000000007;
        }
        return cnt%1000000007;
    }
};

저 부등호와 right, left의 움직임이 많이 헷갈렸고, 마지막에 cnt를 더해줄 때의 min max써주는 것들도 좀 헷갈렸다.

그림으로 보자면 이렇게 볼 수 있다. lower bound l(3) / upper bound r(7)을 구하면 그 중간 5를 생략하고 count값을 구할 수 있다. lower -upper간 거리가 멀수록 선형탐색보다 빠른 속도를 낼 수 있다.

 

Python 풀이

from bisect import *
class Solution:
    def waysToSplit(self, nums: List[int]) -> int:
        mod = 10**9+7
        
        pre_sum = [0]
        for num in nums:
            pre_sum.append(pre_sum[-1]+num)
            
        res = 0
        
        for i in range(1, len(nums)-1):
            l = bisect_left(pre_sum, 2*pre_sum[i])
            r = bisect_right(pre_sum, (pre_sum[i]+pre_sum[-1])//2)
            res += max(0, min(len(nums), r) - max(i+1, l))

        return res%mod