Posts Partitions (InterviewBit)
Post
Cancel

Partitions (InterviewBit)

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])

interviewbit

SOLUTION

BASIC APPROACH (ACCEPTED)

If the totalSum/3 is not divisible by 3, we can simply return 0 since there is no way to divide it equally into three parts with equal sum.

Start iterating from left to right and track the sum assuming it to be for the first partition. Once we have this sum as totalSum/3, use another loop to track the sum of second partition. This partition will obviously start after the previous one ended.

Once the sum of second partition also becomes totalSum/3, we just have to verify if there is at least 1 element left to form the third partition.

We do not really need to calculate the sum of third partition because it has to be totalSum/3 (We have already made sure that first two partitions have sum totalSum/3 and we know totalSum is a multiple of 3).

Finally, if there is at least 1 element left to form the third partition, increase the count by 1.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
public class Solution {

    public int solve(int A, int[] B) {

        int n = B.length;

        int totalSum = getTotalSum(B);

        //the total sum should be divisible by 3 so that 3 partitions with equal sum can be made
        if(totalSum%3 != 0) return 0;

        //ans: total num of ways to create three partitions with equal sum
        int count = 0;

        //each partition must have sum as totalSum/3
        int neededPartitionSum = totalSum/3;

        //sum of first partition
        int firstPartitionSum = 0;

        for(int i=0; i<n; i++){

            firstPartitionSum += B[i];

            //if sum equals totalSum/3, first partition found
            if(firstPartitionSum == neededPartitionSum){

                //sum of second partition
                int secondPartitionSum = 0;

                //start from index after the end of first partition
                for(int j=i+1; j<n; j++){

                    secondPartitionSum += B[j];

                    //second partition found
                    if(secondPartitionSum == neededPartitionSum){

                        //the sum of third partition will be
                        //totalSum - sum of first partition - sum of third partition
                        //= 3x - x - x = x
                        //edge case: first partition and second partition use all the elements in the array
                        //so, the third partition cannot be created

                        //if there are more elements using which 3rd partition can be made, increase the count
                        if(j < n-1) count++;


                    }

                }


            }

        }

        return count;


    }

    public int getTotalSum(int[] A){
        int s = 0;
        for(int i=0; i<A.length; i++) s += A[i];

        return s;
    }

}

BETTER APPROACH

We have to form three partitions in an array: (partition one)(partition two)(partition three)

We know that sum of each partition will be totalSum/3. If totalSum is not divisible by 3, simply return 0.

We can find ALL END INDICES of the FIRST PARTITION by using totalSum/3 value, by iterating from left to right.

We can find ALL START INDICES of the THIRD PARTITION by using totalSum/3 value, by iterating from right to left.

Now, the only thing left is to check if the PARTITION TWO which can be formed using any of the combination from PARTITION ONE end index and PARTITION THREE start index is valid (its sum should also be totalSum/3).

For every possible end index of partition one, say i, we will initialize the start of second partition which will be i+1.

For every possible start index of partition three, say j, we will initialize the end of second partition which will be j-1.

(partition one)(partition two)(partition three) ->(0, i)(i+1, j-1)(j,n-1)

Now the only part left is to verify if the PARTITION TWO from i+1 to j-1 has the sum as totalSum/3, for which we can use our prefix sum array. If it is, increase the count by 1.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
public class Solution {

    public int solve(int A, int[] B) {

        int n = B.length;

        // Calculate the total sum of the array
        int totalSum = getTotalSum(B);

        // If the total sum cannot be evenly divided into 3 parts, return 0
        if(totalSum % 3 != 0) return 0;

        // Calculate the prefix sum array of the input array B
        int[] pf = getPrefixSumArray(B);

        // Find the end indices for the first partition where the sum equals totalSum / 3
        List<Integer> findEndIndexForFirstPartition = findEndIndexForFirstPartition(B, totalSum / 3);

        // Find the start indices for the third partition where the sum equals totalSum / 3
        List<Integer> findStartIndexForThirdPartition = findStartIndexForThirdPartition(B, totalSum / 3);

        int count = 0;

        for(int i = 0; i < findEndIndexForFirstPartition.size(); i++) {

            for(int j = 0; j < findStartIndexForThirdPartition.size(); j++) {

                // There should be at least 1 element between the first and third partitions
                if(findStartIndexForThirdPartition.get(j) - findEndIndexForFirstPartition.get(i) >= 2) {

                    int startOfSecondPartition = findEndIndexForFirstPartition.get(i) + 1;
                    int endOfSecondPartition = findStartIndexForThirdPartition.get(j) - 1;

                    // Calculate the sum of the second partition using prefix sum array
                    int secondPartitionSum = rangeSum(startOfSecondPartition, endOfSecondPartition, pf);

                    // If the sum of the second partition matches the target sum, increment count
                    if(secondPartitionSum == totalSum / 3) {
                        count++;
                    }

                }
            }
        }

        return count;
    }

    // Function to calculate the total sum of an array
    public int getTotalSum(int[] A) {
        int s = 0;

        for(int i = 0; i < A.length; i++) {
            s += A[i];
        }

        return s;
    }

    // Function to find end indices for the first partition with the needed sum
    public List<Integer> findEndIndexForFirstPartition(int[] A, int neededSum) {
        int n = A.length;
        List<Integer> list = new ArrayList<>();

        int sum = 0;

        for(int i = 0; i <= n - 3; i++) {
            sum += A[i];
            if(sum == neededSum) {
                list.add(i);
            }
        }

        return list;
    }

    // Function to find start indices for the third partition with the needed sum
    public List<Integer> findStartIndexForThirdPartition(int[] A, int neededSum) {
        int n = A.length;
        List<Integer> list = new ArrayList<>();

        int sum = 0;

        for(int i = n - 1; i >= 2; i--) {
            sum += A[i];
            if(sum == neededSum) {
                list.add(i);
            }
        }

        return list;
    }

    // Function to generate the prefix sum array of an input array
    public int[] getPrefixSumArray(int[] A) {
        int n = A.length;
        int[] pf = new int[n];

        pf[0] = 0;
        for(int i = 1; i < n; i++) {
            pf[i] = pf[i - 1] + A[i];
        }

        return pf;
    }

    // Function to calculate the sum of elements within a range using the prefix sum array
    public int rangeSum(int start, int end, int[] pf) {
        if(start == 0) return pf[end];
        return pf[end] - pf[start - 1];
    }
}

FURTHER OPTIMIZATION

In the previous approach, while calculating the end index of first partition and start index of the third partition, we have already made sure that their sum should be totalSum/3 each. We have also verified at the start that the total sum is a multiple of 3. So, the left partition, which is the second partition has to be of sum totalSum/3. We don’t need to really calculate it again. We can just check if there is at least 1 element which can be used to form the second partition. In that case, we can increase the count by 1.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
public class Solution {

    public int solve(int A, int[] B) {

        int n = B.length;

        int totalSum = getTotalSum(B);

        if(totalSum%3 != 0) return 0;

        List<Integer> findEndIndexForFirstPartition = findEndIndexForFirstPartition(B, totalSum/3);
        List<Integer> findStartIndexForThirdPartition = findStartIndexForThirdPartition(B, totalSum/3);

        int count = 0;

        for(int i=0; i<findEndIndexForFirstPartition.size(); i++){

            for(int j=0; j<findStartIndexForThirdPartition.size(); j++){

                //there should be at least 1 element between first and third partition
                //the first and third partitions already have sum as totalSum/3
                //so the second partition, if possible, must have the needed sum
                if(findStartIndexForThirdPartition.get(j) - findEndIndexForFirstPartition.get(i) >= 2){
                    count++;
                }

            }

        }

        return count;

    }

    public int getTotalSum(int[] A){
        int s = 0;

        for(int i=0; i<A.length; i++) s += A[i];

        return s;
    }

    public List<Integer> findEndIndexForFirstPartition(int[] A, int neededSum){

        int n = A.length;
        List<Integer> list = new ArrayList<>();

        int sum = 0;

        for(int i=0; i<=n-3; i++){
            sum += A[i];
            if(sum == neededSum){
                list.add(i);
            }
        }

        return list;

    }

    public List<Integer> findStartIndexForThirdPartition(int[] A, int neededSum){

        int n = A.length;

        List<Integer> list = new ArrayList<>();

        int sum = 0;

        for(int i=n-1; i>=2; i--){
            sum += A[i];
            if(sum == neededSum){
                list.add(i);
            }
        }

        return list;

    }

}
This post is licensed under CC BY 4.0 by the author.