Partitions

Problem Description

You are given a 1D integer array B containing A integers.

Count the number of ways to split all the elements of the array into 3 contiguous parts so that the sum of elements in each part is the same.

Such that : sum(B[1],..B[i]) = sum(B[i+1],...B[j]) = sum(B[j+1],...B[n])

Problem Constraints

1 <= A <= 10^5
-10^9 <= B[i] <= 10^9

Input Format

First argument is an integer A.
Second argument is an 1D integer array B of size A.

Output Format

Return a single integer denoting the number of ways to split the array B into three parts with the same sum.

Example Input

Input 1:
 A = 5
 B = [1, 2, 3, 0, 3]

Input 2:
 A = 4
 B = [0, 1, -1, 0]

Example Output

Output 1:
 2

Output 2:
 1

Example Explanation

Explanation 1:
 There are no 2 ways to make partitions -
 1. (1,2)+(3)+(0,3).
 2. (1,2)+(3,0)+(3).

Explanation 2:
 There is only 1 way to make partition -
 1. (0)+(-1,1)+(0).

Solution

import Foundation

class Solution {
    
    func findCount(_ A: inout [Int], _ B:Int) -> Int {
        var l = 0
        var h = A.count - 1
        while l <= h {
            let mid = (l+h) / 2
            let v = A[mid]
            if v == B {
                return (A.count - mid)
            } else if v > B {
                h = mid - 1
            } else {
                l = mid + 1
            }
        }
        return A.count - l 
    }
    
	func solve(_ A: inout Int, _ B: inout [Int]) -> Int {
        let total = B.reduce(0,+)
        if total % 3 != 0 {
            return 0
        }
        let part = total / 3
        
        var pre = [Int]()
        // var suf = [Int]()
        var suf = [Int](repeating:0,count:B.count)
        
        let len = A
        
        var accum = 0 
        for i in 0..<len {
            accum += B[i]
            if accum == part {
                pre.append(i)
            }
        }
        
        // accum = 0 
        // for i in (0..<len).reversed(){
        //      accum += B[i]
        //      if accum == part {
        //         suf.append(i)
        //     }
        // }
        
        // suf.reverse()
        
        accum = 0
        var sufCount = 0
        for i in (0..<len).reversed(){
            accum += B[i]
            if accum == part {
                sufCount += 1
            }
            suf[i] = sufCount
        }
        
        
        
        var result = 0
            
        for i in pre {
            // if pre[i] == 0 {
            //     continue
            // }   
            var c = 0
            if (i+2) < len {
            // for j in (i+2)..<len {
            //     if suf[j] > 0 {
            //         c += 1
            //     }
            // }    
                // c =  findCount(&suf,i+2)  
                c = suf[i+2]         
            } 
            
            if c == 0 {
                break
            }

            // let c = Array(suf[i+2..<len]).reduce
            
            result += c
            
        }
        return result
        
        
        
        // if part == 0 {
        //     var zeroParts = [(Int,Int)]()
        //     var accum = 0
        //     var start = 0
        //     for i in 0..<A {
        //         accum += B[i]
        //         if accum == 0 {
        //             zeroParts.append((start,i))     
        //             start = i + 1
        //         }
        //     }
        //     if zeroParts.count < 3 {
        //         return 0
        //     }
            
        //     var result = 0
        //     let len = zeroParts.count-2
        //     for i in 0..<len {
        //         // let l = i
        //         let r = len - i - 1
        //         result += (r+1)
                
        //     }
            
        //     return result
        // }
        
        // var zeroParts12 = [(Int,Int)]()
        // var zeroParts23 = [(Int,Int)]()
        // var equalParts = [(Int,Int)]()
        
        // // let len = A
        // var accum = 0
        // var start = 0
        // for i in 0..<A {
        //     accum += B[i]
        //     if accum == 0 {
        //         if equalParts.count == 1 {
        //             zeroParts12.append((start,i))
        //         } else if equalParts.count == 2 {
        //             zeroParts23.append((start,i))
        //         }       
        //         start = i + 1
        //     } else if accum == part {
        //         equalParts.append((start,i))
        //         start = i + 1
        //         accum = 0
        //     } else if accum > part {
        //         return 0
        //     }
        // }
        
        // if equalParts.count < 3 {
        //     return 0
        // }
        
        // return (zeroParts12.count+1)*(zeroParts23.count+1)
        
	}
}

References

• InterviewBit > Partitions