2736. Maximum Sum Queries

HardStackBinary Indexed TreeSegment TreeArrayBinary SearchSortingMonotonic Stack
Leetcode Link

Problem Description

In this LeetCode problem, we are given two arrays nums1 and nums2 of the same length n, and another 2D array queries where each query is represented as a pair [x, y]. The task is to process each query by looking for indices j such that nums1[j] >= x and nums2[j] >= y. For each query, we want to find the maximum sum of nums1[j] + nums2[j] over all indices j that meet the criteria. If no such j exists for a query, we should return -1 for that query. The output should be an array containing the result for each query in the same order as the queries were presented.

Intuition

To solve this problem, we can use several steps:

  1. Preprocessing: Combine nums1 and nums2 into a new array where each element is a pair [nums1[j], nums2[j]]. This allows us to treat the two conditions in the queries (comparing to x and y) more easily as we can sort and traverse this new array based on these conditions.

  2. Sorting: We sort our newly constructed array by the first elements of the pairs (values from nums1). Likewise, we sort the queries array by the x value. This sorting step allows us to efficiently traverse the arrays while respecting the constraints of each query (x and y conditions).

  3. Using TreeMap: We use a TreeMap where keys represent values from nums2 and values in the map represent the max sum nums1[j] + nums2[j] for a given nums2[j]. TreeMap is chosen because it provides logarithmic time cost for operations like searching for the smallest key greater than or equal to a given key (ceilingKey) or the greatest key less than or equal to a given key (floorKey). These operations are crucial for efficiently answering the queries with respect to the y condition.

  4. Processing Queries in Reverse: We start from the last query after sorting and move backwards. This is because we have sorted the queries by x. As we go through each pair in nums1 and nums2 starting from the highest nums1[j], we update our map accordingly. For each query, as x becomes smaller or equal to the current nums1[j], we know that all future queries will also satisfy this nums1[j] condition (due to sorted queries).

  5. Updating TreeMap and Results: When a nums1[j] meets the condition for the current query, we update the TreeMap if necessary (i.e., if we encounter a higher nums2[j] than previously recorded). We do this by removing entries in the TreeMap that are no longer relevant - when they do not represent the maximum sum for their nums2[j]. Finally, for each query, we seek the smallest key in TreeMap (ceilingKey) greater than or equal to y to satisfy the second condition of the query. If such a key exists, we have our answer for that query; otherwise, we return -1.

Overall, the solution finds an efficient way to keep track of the maximum sums while honoring both conditions in each query through sorted traversal and data structure choice.

Learn more about Stack, Segment Tree, Binary Search, Sorting and Monotonic Stack patterns.

Solution Approach

The implementation given in the Reference Solution Approach follows these steps:

  1. Initialization and Preprocessing: We create an array a where each element is a pair [nums1[i], nums2[i]]. Similarly, we create an array b for the queries where each element [b[i][0], b[i][1], b[i][2]] represents the x, y, and index of the query.

  2. Sorting Arrays: We sort the array a according to the nums1 values (first elements of the pairs) and the array b according to the x value of the queries:

    Arrays.sort(a, (o1, o2) -> o1[0] - o2[0]);
Arrays.sort(b, (o1, o2) -> o1[0] - o2[0]);

This sorting lets us iterate over the arrays while respecting the order imposed by the queries' x values and the pairs' nums1 values.

  1. Using TreeMap: The TreeMap named map is used to store the maximum sums. Keys in the map are nums2 values, and values are the maximum sum nums1[j] + nums2[j] encountered for that key.

  2. Processing Queries in Reverse Order: By traversing the queries in reverse (starting from the largest x), we handle the queries with higher x values first. This is essential as it aligns with our sorted array of pairs a.

    for (int i = m - 1, j = n - 1; i >= 0; i--) {
    int x = b[i][0], y = b[i][1], idx = b[i][2];
    // Rest of the loop
}

As we encounter elements in a that satisfy the nums1[j] >= x condition, we:

  1. Update Max and TreeMap: We keep track of a variable max which stores the highest nums2 value encountered that has satisfied the nums1[j] >= x condition for the current query. If a new max is found, we update the TreeMap by removing keys that represent sums smaller than the current max sum for their corresponding nums2 value:

    if (max < a[j][1]) {
    // Updating the max and TreeMap
}

With map.ceilingKey(y), we find the smallest nums2 value in the map that is greater than or equal to y, satisfying the queries' y value condition.

  1. Populating the Results: For each query, we use the index stored in b[i][2] and assign to res[idx] either the value of map.get(key) in case a key was found, or -1 if no such key exists:

    if (key == null)
    res[idx] = -1;
else
    res[idx] = map.get(key);

Finally, we return the res array as the solution, containing the answers to all the queries as specified by the problem.

By combining sorting, binary search features of TreeMap, and careful updates to the data structure, the algorithm efficiently satisfies the complex query constraints in an optimized manner.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's illustrate the solution approach with a small example:

Suppose nums1 = [3, 5, 1], nums2 = [2, 4, 3], and our queries are [[4, 2], [2, 3], [3, 1]]. We want to process these queries to find the maximum sum of pairs from nums1 and nums2 that meet the query conditions on x and y.

  1. Initialization and Preprocessing: We create an array a with paired elements, resulting in a = [[3, 2], [5, 4], [1, 3]]. Then, we prepare the queries by including the index of each query in array b, thus b = [[4, 2, 0], [2, 3, 1], [3, 1, 2]].

  2. Sorting Arrays: Array a after sorting by nums1 values becomes [[1, 3], [3, 2], [5, 4]], and array b sorted by x values becomes [[2, 3, 1], [3, 1, 2], [4, 2, 0]].

  3. Using TreeMap: We initiate an empty TreeMap named map.

  4. Processing Queries in Reverse Order: We start from the last query [4, 2, 0] since it has the largest x.

  5. Update Max and TreeMap: While traversing a in reverse, we find the first pair [5, 4] satisfies nums1[j] >= 4, so we add it to the map with map.put(4, 5 + 4) unless there's already a larger sum for key 4.

  6. Populating the Results: For the query [4, 2, 0], we now look for the smallest key in the map that is greater than or equal to y = 2, which is key 4, and the current max sum is 5 + 4 = 9. Hence, we put 9 in the result array at index 0: res[0] = 9.

Next, we move to the second last query [3, 1, 2]. Our map currently contains the sum for nums2[j] = 4 but since nums1[j] >= 3 is also satisfied by [3, 2] which is the second pair in a, we check if that gives us a higher sum. However, we already have a larger sum in the map, so we continue.

For query [3, 1, 2], we find the smallest key in the map greater than or equal to 1, which is 4, and the corresponding sum is still 9. So res[2] = 9.

Finally, we look at query [2, 3, 1]. Since all elements in a satisfy nums1[j] >= 2, we don't need to add anything new to the map. Now, we look for the smallest key greater than or equal to 3, which does not exist in the map, so res[1] = -1.

The final results array is res = [9, -1, 9], which is the solution we're looking for.

In this example, we can see that by pre-sorting the arrays and using a map to keep track of the maximum sums, the solution methodically processes each query according to the required conditions and efficiently finds the highest sum that satisfies each query.

Solution Implementation

1from sortedcontainers import SortedDict
2
3class Solution:
4
5    def maximum_sum_queries(self, nums1, nums2, queries):
6        n = len(nums1)  # Number of elements in nums1 and nums2.
7        m = len(queries)  # Number of queries.
8      
9        # Create a list 'pairs' to store values and their corresponding weights from nums1 and nums2.
10        pairs = [(nums1[i], nums2[i]) for i in range(n)]
11      
12        # Create a list 'sorted_queries' to store queries along with their original indices for result mapping.
13        sorted_queries = [(queries[i][0], queries[i][1], i) for i in range(m)]
14      
15        # Sort pairs by value in ascending order.
16        pairs.sort(key=lambda x: x[0])
17        # Sort sorted_queries by starting value in ascending order.
18        sorted_queries.sort(key=lambda x: x[0])
19
20        # SortedDict to maintain weights and corresponding maximum possible sum.
21        weight_to_sum_map = SortedDict()
22        result = [-1] * m  # Initialize result array for storing final results with -1.
23        max_weight = -1  # Used to keep track of the maximum weight encountered so far.
24
25        # Iterate over the sorted_queries in reverse.
26        j = n - 1
27        for start_value, end_weight, query_index in reversed(sorted_queries):
28
29            # Keep evaluating pairs till the values are within the query's start value.
30            while j >= 0 and pairs[j][0] >= start_value:
31                value, weight = pairs[j]
32                if max_weight < weight:
33                    max_weight = weight
34                    # Remove entries in the map where the sum of value and weight is not greater
35                    # than the current pair's sum.
36                    keys_to_remove = [key for key in weight_to_sum_map
37                                      if weight_to_sum_map[key] <= value + weight]
38                    for key in keys_to_remove:
39                        del weight_to_sum_map[key]
40                    # Put the new max weight and corresponding sum into the map.
41                    weight_to_sum_map[max_weight] = value + weight
42                j -= 1
43          
44            # Find the least weight that is greater than or equal to end_weight.
45            key = weight_to_sum_map.bisect_left(end_weight)
46            if key != len(weight_to_sum_map):
47                result[query_index] = weight_to_sum_map.peekitem(key)[1]
48      
49        return result  # Return the final results array.
50
1import java.util.Arrays;
2import java.util.TreeMap;
3
4class Solution {
5
6    public int[] maximumSumQueries(int[] nums1, int[] nums2, int[][] queries) {
7        int n = nums1.length; // Number of elements in nums1 and nums2.
8        int m = queries.length; // Number of queries.
9      
10        // Create an array 'pairs' to store values and their corresponding weights from nums1 and nums2.
11        int[][] pairs = new int[n][2];
12        for (int i = 0; i < n; i++) {
13            pairs[i][0] = nums1[i]; // Value
14            pairs[i][1] = nums2[i]; // Weight
15        }
16
17        // Create an array 'sortedQueries' to store queries along with their original indices for result mapping.
18        int[][] sortedQueries = new int[m][3];
19        for (int i = 0; i < m; i++) {
20            sortedQueries[i][0] = queries[i][0]; // Starting value
21            sortedQueries[i][1] = queries[i][1]; // Ending weight
22            sortedQueries[i][2] = i;             // Original index of query
23        }
24
25        // Sort pairs by value in ascending order.
26        Arrays.sort(pairs, (o1, o2) -> o1[0] - o2[0]);
27        // Sort sortedQueries by starting value in ascending order.
28        Arrays.sort(sortedQueries, (o1, o2) -> o1[0] - o2[0]);
29
30        // TreeMap to maintain weights and corresponding maximum possible sum.
31        TreeMap<Integer, Integer> map = new TreeMap<>();
32        int[] result = new int[m]; // Result array for storing final results.
33        int max = -1; // Used to keep track of the maximum weight encountered so far.
34
35        // Iterate over the sortedQueries in reverse.
36        for (int i = m - 1, j = n - 1; i >= 0; i--) {
37            int startValue = sortedQueries[i][0], endWeight = sortedQueries[i][1];
38            int queryIndex = sortedQueries[i][2];
39
40            // Keep evaluating pairs till the values are within query's start value.
41            while (j >= 0 && pairs[j][0] >= startValue) {
42                if (max < pairs[j][1]) {
43                    max = pairs[j][1];
44                    Integer key = map.floorKey(max);
45                  
46                    // Remove entries in the map where the sum of value and weight is not greater
47                    // than the current pair's sum.
48                    while (key != null && map.get(key) <= pairs[j][0] + pairs[j][1]) {
49                        map.remove(key);
50                        key = map.floorKey(key);
51                    }
52                    // Put the new max value and corresponding sum into the map.
53                    map.put(max, pairs[j][0] + pairs[j][1]);
54                }
55                j--;
56            }
57          
58            Integer key = map.ceilingKey(endWeight);
59            // If there is no weight that can handle the query, set result as -1.
60            if (key == null) {
61                result[queryIndex] = -1;
62            } else {
63                // Otherwise, set the result to the corresponding sum.
64                result[queryIndex] = map.get(key);
65            }
66        }
67        return result; // Return the final results array.
68    }
69}
70
1#include <vector>
2#include <map>
3#include <algorithm>
4
5class Solution {
6public:
7    // Function to calculate the maximum sum for the weight intervals specified in queries.
8    std::vector<int> maximumSumQueries(std::vector<int>& nums1, std::vector<int>& nums2, std::vector<std::vector<int>>& queries) {
9        int n = nums1.size();  // Number of elements in nums1 and nums2.
10        int m = queries.size();  // Number of queries.
11
12        // Create a vector of pairs to store values and their corresponding weights from nums1 and nums2.
13        std::vector<std::pair<int, int>> pairs(n);
14        for (int i = 0; i < n; ++i) {
15            pairs[i] = {nums1[i], nums2[i]}; // Pair(value, weight)
16        }
17
18        // Create a vector to store queries along with their original indices for result mapping.
19        std::vector<std::tuple<int, int, int>> sortedQueries(m);
20        for (int i = 0; i < m; ++i) {
21            sortedQueries[i] = std::make_tuple(queries[i][0], queries[i][1], i); // Tuple(startValue, endWeight, queryIndex)
22        }
23
24        // Sort pairs by value in ascending order.
25        std::sort(pairs.begin(), pairs.end());
26        // Sort sortedQueries by starting value in ascending order.
27        std::sort(sortedQueries.begin(), sortedQueries.end(), 
28            [](const std::tuple<int, int, int>& a, const std::tuple<int, int, int>& b) {
29                return std::get<0>(a) < std::get<0>(b);
30            });
31
32        // TreeMap to maintain weights and corresponding maximum possible sum.
33        std::map<int, int> map;
34        std::vector<int> result(m); // Result vector for storing final results.
35        int maxWeight = -1;  // Used to keep track of the maximum weight encountered so far.
36
37        // Iterate over the sortedQueries in reverse.
38        for (int i = m - 1, j = n - 1; i >= 0; --i) {
39            int startValue = std::get<0>(sortedQueries[i]), endWeight = std::get<1>(sortedQueries[i]);
40            int queryIndex = std::get<2>(sortedQueries[i]);
41
42            // Keep evaluating pairs till the values are within the query's start value.
43            while (j >= 0 && pairs[j].first >= startValue) {
44                if (maxWeight < pairs[j].second) {
45                    maxWeight = pairs[j].second;
46                    auto it = map.lower_bound(maxWeight);
47                  
48                    // Remove entries in the map where the sum of value and weight is not greater than the current pair's sum.
49                    while (it != map.begin() && (--it)->second <= pairs[j].first + pairs[j].second) {
50                        it = map.erase(it);
51                    }
52                    // Put the new max weight and corresponding sum into the map.
53                    map[maxWeight] = pairs[j].first + pairs[j].second;
54                }
55                --j;
56            }
57          
58            // Find the smallest key in map that is greater than or equal to endWeight.
59            auto it = map.lower_bound(endWeight);
60            // If there is no weight that can handle the query, set result as -1.
61            if (it == map.end()) {
62                result[queryIndex] = -1;
63            } else {
64                // Otherwise, set the result to the corresponding sum.
65                result[queryIndex] = it->second;
66            }
67        }
68        return result;  // Return the final results vector.
69    }
70};
71
1function maximumSumQueries(nums1: number[], nums2: number[], queries: number[][]): number[] {
2    const n: number = nums1.length; // Number of elements in nums1 and nums2.
3    const m: number = queries.length; // Number of queries.
4  
5    // Create an array 'pairs' to store values and their corresponding weights from nums1 and nums2.
6    const pairs: number[][] = new Array(n).fill(0).map(() => Array(2));
7    for (let i = 0; i < n; i++) {
8        pairs[i][0] = nums1[i]; // Value
9        pairs[i][1] = nums2[i]; // Weight
10    }
11
12    // Create an array 'sortedQueries' to store queries along with their original indices for result mapping.
13    const sortedQueries: number[][] = new Array(m).fill(0).map(() => Array(3));
14    for (let i = 0; i < m; i++) {
15        sortedQueries[i][0] = queries[i][0]; // Starting value
16        sortedQueries[i][1] = queries[i][1]; // Ending weight
17        sortedQueries[i][2] = i;             // Original index of query
18    }
19
20    // Sort pairs by value in ascending order and sortedQueries by starting value in ascending order.
21    pairs.sort((a, b) => a[0] - b[0]);
22    sortedQueries.sort((a, b) => a[0] - b[0]);
23
24    // TreeMap equivalent in TypeScript using a Map to maintain weights and corresponding maximum possible sum.
25    const map: Map<number, number> = new Map<number, number>();
26    const result: number[] = new Array(m).fill(-1); // Result array for storing final results initialized with -1.
27    let max: number = -1; // Used to keep track of the maximum weight encountered so far.
28
29    // Iterate over the sortedQueries in reverse.
30    for (let i = m - 1, j = n - 1; i >= 0; i--) {
31        const startValue: number = sortedQueries[i][0];
32        const endWeight: number = sortedQueries[i][1];
33        const queryIndex: number = sortedQueries[i][2];
34
35        // Keep evaluating pairs till the values are within query's start value.
36        while (j >= 0 && pairs[j][0] >= startValue) {
37            if (max < pairs[j][1]) {
38                max = pairs[j][1];
39                let key: number | undefined = Array.from(map.keys()).reverse().find(k => k <= max);
40              
41                // Remove entries in the map where the sum of value and weight is not greater
42                // than the current pair's sum.
43                while (key !== undefined && map.get(key)! <= pairs[j][0] + pairs[j][1]) {
44                    map.delete(key);
45                    key = Array.from(map.keys()).reverse().find(k => k <= max);
46                }
47                // Put the new max value and corresponding sum into the map.
48                map.set(max, pairs[j][0] + pairs[j][1]);
49            }
50            j--;
51        }
52      
53        const key: number | undefined = Array.from(map.keys()).find(k => k >= endWeight);
54        // If there is no weight that can handle the query, keep the result as -1 (already initialized).
55        if (key !== undefined) {
56            // Otherwise, set the result to the corresponding sum.
57            result[queryIndex] = map.get(key)!;
58        }
59    }
60    return result; // Return the final results array.
61}
62
63// Example usage:
64// const nums1 = [1, 2, 3];
65// const nums2 = [3, 2, 1];
66// const queries = [[1, 2], [2, 3]];
67// console.log(maximumSumQueries(nums1, nums2, queries));
68

Time and Space Complexity

Time Complexity

The time complexity of the main operations in the code are:

  1. Sorting the array a: Sorting an array of n elements has a time complexity of O(n log n).

  2. Sorting the array b: Sorting an array of m elements also has a time complexity of O(m log m).

  3. The while loop that processes arrays a and b:

    • In the worst case, for each query in array b, it processes elements in array a once. Since b is iterated in reverse, and the index j is only decremented, array a is effectively processed only once overall regardless of the number of queries.
    • The operations inside the while loop have to do with TreeMap operations. The worst-case time complexities for floorKey, ceilingKey, and remove methods in a TreeMap are O(log k), where k is the number of elements in the TreeMap.

  4. The overall time complexity thus becomes O(n log n) + O(m log m) + O(m * log k). Since n could be at most the same as m, we can consider n and m to be interchangeable from a big O notation standpoint. Therefore, the overall time complexity simplifies to O(n log n) + O(m * log k).

  5. Note that k could at worst be n, since you're inserting at most n different sums into the TreeMap. So, the time complexity of the TreeMap operations could be seen as O(log n).

Thus, the time complexity of the code is O(n log n + m log n).

Space Complexity

The space complexity is determined by:

  1. The space used to store arrays a and b, which will use O(n) and O(m) space respectively.

  2. The space used by the TreeMap map, which in the worst case will store n entries, thus O(n).

  3. The additional space for auxiliary variables and the output array res, which is O(m).

Adding these up, the space complexity is O(n + m) because they are not nested structures and can be considered as summing up linearly.

Learn more about how to find time and space complexity quickly using problem constraints.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:
Question 1 out of 10

What data structure does Breadth-first search typically uses to store intermediate states?

Recommended Readings

Stack Intro

Stack Intro Imagine you have a pile of books on your desk If you want to add a new book you place it on top If you want to read a book you take it from the top And if you simply want to see which book is on the top you

Segment Tree

Segment Tree Faster Range Queries For this article we want to introduce the idea of a Segment Tree Segment Trees allow us to quickly perform range queries as well as range updates Suppose we have an array and we want to know the sum of a particular range of numbers as well

Binary Search Speedrun

Binary Search Speedrun For each of the Speedrun questions you will be given a binary search related problem and a corresponding multiple choice question The multiple choice questions are related to the techniques and template s introduced in the binary search section It's recommended that you have gone through at

Want a Structured Path to Master System Design Too? Don’t Miss This!