Skip to main content

Prefix Sum

Can You Eat Your Favorite Candy on Your Favorite Day?

Solution

Java Code

class Solution {
    public boolean[] canEat(int[] candiesCount, int[][] queries) {
        long[] candiesPrefixCount = new long[candiesCount.length];
        candiesPrefixCount[0] = candiesCount[0];
        for(int i = 1; i < candiesCount.length; i++) {
            candiesPrefixCount[i] = candiesPrefixCount[i - 1] + candiesCount[i];
        }

        boolean[] answer = new boolean[queries.length];
        for(int i = 0; i < queries.length; i++) {
            int[] query = queries[i];
            int favCandy = query[0];
            long favDay = query[1];
            long dailyCap = query[2];

            long totalCandiesToBeEatenBefore = (favCandy == 0) ? 0 : candiesPrefixCount[favCandy - 1];

            if(totalCandiesToBeEatenBefore >= (favDay + 1) * dailyCap) {
                answer[i] = false;
            } else if((favDay + 1) > candiesPrefixCount[favCandy]) {
                answer[i] = false;
            } else { 
                answer[i] = true;
            }
        }

        return answer;
    }
}

C++ Code

class Solution {
public:
    vector<bool> canEat(vector<int>& candiesCount, vector<vector<int>>& queries) {
        int n = candiesCount.size();
        vector<long> candiesPrefixCount(n);
        candiesPrefixCount[0] = candiesCount[0];
        for (int i = 1; i < n; i++) {
            candiesPrefixCount[i] = candiesPrefixCount[i - 1] + candiesCount[i];
        }

        vector<bool> answer;
        for (const auto& query : queries) {
            int favCandy = query[0];
            long favDay = query[1];
            long dailyCap = query[2];

            long totalCandiesToBeEatenBefore = (favCandy == 0) ? 0 : candiesPrefixCount[favCandy - 1];

            if (totalCandiesToBeEatenBefore >= (favDay + 1) * dailyCap) {
                answer.push_back(false);
            } else if ((favDay + 1) > candiesPrefixCount[favCandy]) {
                answer.push_back(false);
            } else {
                answer.push_back(true);
            }
        }

        return answer;
    }
};

Python Code

class Solution:
    def canEat(self, candiesCount: List[int], queries: List[List[int]]) -> List[bool]:
        n = len(candiesCount)
        candiesPrefixCount = [0] * n
        candiesPrefixCount[0] = candiesCount[0]
        for i in range(1, n):
            candiesPrefixCount[i] = candiesPrefixCount[i - 1] + candiesCount[i]

        answer = []
        for query in queries:
            favCandy, favDay, dailyCap = query

            totalCandiesToBeEatenBefore = 0 if favCandy == 0 else candiesPrefixCount[favCandy - 1]

            if totalCandiesToBeEatenBefore >= (favDay + 1) * dailyCap:
                answer.append(False)
            elif (favDay + 1) > candiesPrefixCount[favCandy]:
                answer.append(False)
            else:
                answer.append(True)

        return answer

Complexity Analysis

Time Complexity: O(N + Q), where N is the length of the given array candiesCount and Q is the number of queries queries.
Space Complexity: O(N), for building candiesPrefixCount array. We will not consider answer array in the space complexity as this is required in the output.