Skip to main content

Prefix Sum

Corporate Flight Bookings

Solution 1: Using Two Arrays

In the previous problems, we have seen that we need one extra index to update the right range. But in this question, we need to return exactly n size array.

Therefore, we can do this using two different arrays.

  1. Using tempPrefixSum which is as usual in the previous problems i.e. one extra size.
  2. Finally, creating a new array prefixSum to ignore the last index because we anyway don't need it.
💡
One important point: In a normal scenario of such questions we have seen that whenever we are given range updates i.e. [left, right]. We increment at prefixSum[left] and decrement at prefixSum[right + 1].

But in this question the range is not 0 index. For example, consider the booking [2, 5] which means the following:

We need to update prefixSum[2 - 1] because 2 means 2nd flight which is represented by 1st index in prefixSum.
Similarly: we need to update prefixSum[5] because 5 means 5th flight which is represented by 4th index in prefixSum.

Java Code

class Solution {
    public int[] corpFlightBookings(int[][] bookings, int n) {
        int[] tempPrefixSum = new int[n + 1];

        for(int[] booking : bookings) {
            int fromFlight = booking[0];
            int toFlight = booking[1];
            int seatsBooked = booking[2];
            tempPrefixSum[fromFlight - 1] += seatsBooked;
            tempPrefixSum[toFlight] += -seatsBooked;
        }
        
        int[] prefixSum = new int[n];
        prefixSum[0] = tempPrefixSum[0];
        for(int i = 1; i < prefixSum.length; i++) {
            prefixSum[i] = prefixSum[i - 1] + tempPrefixSum[i];
        }
        return prefixSum;
    }
}

C++ Code

class Solution {
public:
    vector<int> corpFlightBookings(vector<vector<int>>& bookings, int n) {
        vector<int> tempPrefixSum(n + 1, 0);

        for(auto &booking : bookings) {
            int fromFlight = booking[0];
            int toFlight = booking[1];
            int seatsBooked = booking[2];
            tempPrefixSum[fromFlight - 1] += seatsBooked;
            tempPrefixSum[toFlight] += -seatsBooked;
        }
        
        vector<int> prefixSum(n, 0);
        prefixSum[0] = tempPrefixSum[0];
        for(int i = 1; i < prefixSum.size(); i++) {
            prefixSum[i] = prefixSum[i - 1] + tempPrefixSum[i];
        }
        return prefixSum;
    }
};

Python Code

class Solution:
    def corpFlightBookings(self, bookings: List[List[int]], n: int) -> List[int]:
        temp_prefix_sum = [0] * (n + 1)
        
        for booking in bookings:
            from_flight = booking[0]
            to_flight = booking[1]
            seats_booked = booking[2]
            temp_prefix_sum[from_flight - 1] += seats_booked
            temp_prefix_sum[to_flight] += -seats_booked
        
        prefix_sum = [0] * n
        prefix_sum[0] = temp_prefix_sum[0]
        for i in range(1, n):
            prefix_sum[i] = prefix_sum[i - 1] + temp_prefix_sum[i]
        return prefix_sum

Complexity Analysis

Time Complexity: O(N + B), where N is the number of flights which is needed to create prefixSum array and B is the number of bookings represented by bookings array.
Space Complexity: O(N), the size of the prefixSum array.

Solution 2: Using only one Array

💡
Consider the range [2, 5]. Do we need to update the last range? No. Because while calculating the prefixSum, we don't really need the last index as it is not useful for us. Hence we can ignore. Just dry and run, you will get it.

Java Code

class Solution {
    public int[] corpFlightBookings(int[][] bookings, int n) {
        int[] prefixSum = new int[n];

        for(int[] booking : bookings) {
            int fromFlight = booking[0];
            int toFlight = booking[1];
            int seatsBooked = booking[2];
            prefixSum[fromFlight - 1] += seatsBooked;
            if(toFlight < n) {
                prefixSum[toFlight] += -seatsBooked;
            }
        }
        
        for(int i = 1; i < prefixSum.length; i++) {
            prefixSum[i] = prefixSum[i - 1] + prefixSum[i];
        }
        return prefixSum;
    }
}

C++ Code

class Solution {
public:
    vector<int> corpFlightBookings(vector<vector<int>>& bookings, int n) {
        vector<int> prefixSum(n, 0);

        for(auto &booking : bookings) {
            int fromFlight = booking[0];
            int toFlight = booking[1];
            int seatsBooked = booking[2];
            prefixSum[fromFlight - 1] += seatsBooked;
            if(toFlight < n) {
                prefixSum[toFlight] += -seatsBooked;
            }
        }
        
        for(int i = 1; i < prefixSum.size(); i++) {
            prefixSum[i] = prefixSum[i - 1] + prefixSum[i];
        }
        return prefixSum;
    }
};

Python Code

class Solution:
    def corpFlightBookings(self, bookings: List[List[int]], n: int) -> List[int]:
        prefix_sum = [0] * n
        
        for booking in bookings:
            from_flight = booking[0]
            to_flight = booking[1]
            seats_booked = booking[2]
            prefix_sum[from_flight - 1] += seats_booked
            if(to_flight < n):
                prefix_sum[to_flight] += -seats_booked
        
        for i in range(1, n):
            prefix_sum[i] = prefix_sum[i - 1] + prefix_sum[i]
        return prefix_sum

Complexity Analysis

Time Complexity: O(N + B), where N is the number of flights which is needed to create prefixSum array and B is the number of bookings represented by bookings array.
Space Complexity: O(N), the size of the prefixSum array but we are using only one array.