Max of All Subarrays of Size k Using Max Heap in C++

This problem, also referred to as the “sliding window maximum” problem, can be efficiently solved using different programming techniques. We’ve already discussed an O(n) time complexity approach using deque in this post – Find maximum of all subarrays of size k

In this post we’ll be solving this problem using a max-heap. Max heap is discussed in detail in these posts:

Problem Statement

Given an array of size n and an integer k, we need to find the maximum array element for each contiguous subarray of size k.

Approach – Using Max Heap

A max heap is a complete binary tree where the value of each node is greater than or equal to the values of its children. This property makes max heap ideal for efficiently retrieving the maximum element. Consider this:

  • Insertion and deletion are logarithmic: Inserting and deleting elements in a max heap take O(log k) time, where k is the size of the heap.
  • Max retrieval is constant: The maximum element can always be found at the root of the heap, allowing for O(1) access.

The idea is to use a max heap to keep track of the maximum values of the current window of size k.

//
//  main.cpp
//  Max of subarrays
//
//  Created on 20/06/24.
//

#include <iostream>
#include <vector>
#include <algorithm>
#include <unordered_map>

using namespace std;

vector<int> maxOfSubarray(vector<int>& nums, int k) {
    vector<int> result;
    vector<pair<int, int>> heap;  // To store value and its index: <value, index>

    // Custom comparator for max heap
    auto cmp = [](const pair<int, int>& a, const pair<int, int>& b) {
        return a.first < b.first;  // Max heap comparator
    };
    
    // Initialize heap with the first k elements
    for (int i = 0; i < k; i++) {
        heap.push_back({nums[i], i});
        push_heap(heap.begin(), heap.end(), cmp);
    }
    
    // Push the maximum element of current window ie 
    // element at the front of heap to result
    result.push_back(heap.front().first);

    for (int i = k; i < nums.size(); i++) {
        
        // Remove the element that goes out of the window
        while (!heap.empty() && heap.front().second <= i - k) {
            pop_heap(heap.begin(), heap.end(), cmp); // Re-heapify
            heap.pop_back(); // Actually remove the element
        }
        
        // Add current element to the heap
        heap.push_back({nums[i], i});
        push_heap(heap.begin(), heap.end(), cmp); // Re-heapify

        // The maximum is always at the front of the max heap
        result.push_back(heap.front().first);
    }

    return result;
}

int main() {
    vector<int> nums = {1, 3, -1, -3, 5, 3, 6, 7};
    int k = 3;
    
    vector<int> result = maxOfSubarray(nums, k);

    cout << "Max elements in subarrays of size " << k << " are: ";
    for (int x : result) {
        cout << x << " ";
    }
    cout << endl;

    return 0;
}

Output

Max elements in subarrays of size 3 are: 3 3 5 5 6 7

Time Complexity: O(n * log n), n is the size of the input vector/array
Auxiliary Space: O(n), for storing heap

Practice Problem: Sliding Window Maximum (LeetCode)

Leave a Reply

Your email address will not be published. Required fields are marked *