Count elements that are greater than at least K elements on its left and right

Given an array arr[] of size n (1 <= n <= 10^5) and a positive integer k, the task is to count all indices i ( 1<= i < n) in the array such that at least k elements to the left of i and at least k elements to the right of i, which are strictly smaller than the value at the ith index (i.e, arr[i]).

Example:

Input: arr = {1,3,6,5,2,1}, k = 2
Output: 2
Explanation: The elements marked in bold are only valid indices, arr = {2,3,6,5,2,3}

Input: arr = {2,2,2}, k = 3
Output: 0

Approach:

The idea is to choose some specific data structure which will keep track of maximum. so that we can check if current index i is greater than maximum element to the left. Lets see about the core idea.

We’ll use Max heap data structure and always maintain its size to be k. For any index i, we’ll check if current value at i (i.e, arr[i]) is greater than the max heap top element, it means there must me at least k elements to the left of i which are less than current value. If current value is less than the max heap’s top then we have to remove the max heap top and insert current element into the heap as this will help us in more reducing the maximum value so it will help other future element while validating. We’ll keep track of this validation in left[] array.

We’ll do the above same from right to left and keep track of validation in right[] array.

Finally, we’ll travel in both array and check if left[i] == 1 or true && right[i] == 1 or true then this is valid index. So, increment our count.

Steps-by-step approach:

• Initialize two priority queues maxH and maxH2 to track the k largest elements, and create vectors for left[] and right[] indices.
• Iterate through the array from left to right, updating the priority queue and marking elements satisfying conditions on the left[].
• Iterate through the array from right to left, updating the second priority queue and marking elements satisfying conditions on the right[].
• Count the number of elements satisfying conditions (left[i] == 1 or true && right[i] == 1 or true) on both sides.
• Return the count as the result.

Below is the implementation of the above approach:

C++

#include<bits/stdc++.h>
using namespace std;

int validElements(vector<int>& nums, int k) {

    // Initialize two priority queues to keep track of the k largest elements
    priority_queue<int> maxH, maxH2;

    int n = nums.size();
    vector<int> left(n), right(n);

    // Iterate over the array from left to right
    for(int i = 0; i < n; i++){
        if(maxH.size() < k){
            maxH.push(nums[i]);
        }
        else{
            if(maxH.top() < nums[i]) left[i] = 1;
            else{
                maxH.pop();
                maxH.push(nums[i]);
            }
        }
    }

    // Iterate over the array from right to left
    for(int i = n - 1; i >= 0; i--){
        if(maxH2.size() < k){
            maxH2.push(nums[i]);
        }
        else{
            if(maxH2.top() < nums[i]) right[i] = 1;
            else{
                maxH2.pop();
                maxH2.push(nums[i]);
            }
        }
    }

    // Count the number of elements that are among the k largest from both sides
    int count = 0;
    for(int i = 0; i < n; i++) if(left[i] && right[i]) count++;

    return count;
}

int main() {

    // input array and k
    vector<int> nums = {2,3,6,5,2,3};
    int k = 2;
    
    // Call the function and print the result
    cout << validElements(nums, k) << endl;
    return 0;
}

Java

import java.util.PriorityQueue;

public class Main {

    public static int validElements(int[] nums, int k) {
        // Initialize two priority queues to keep track of the k largest elements
        PriorityQueue<Integer> maxH = new PriorityQueue<>((a, b) -> b - a);
        PriorityQueue<Integer> maxH2 = new PriorityQueue<>((a, b) -> b - a);

        int n = nums.length;
        int[] left = new int[n];
        int[] right = new int[n];

        // Iterate over the array from left to right
        for (int i = 0; i < n; i++) {
            if (maxH.size() < k) {
                maxH.offer(nums[i]);
            } else {
                if (maxH.peek() < nums[i]) {
                    left[i] = 1;
                } else {
                    maxH.poll();
                    maxH.offer(nums[i]);
                }
            }
        }

        // Iterate over the array from right to left
        for (int i = n - 1; i >= 0; i--) {
            if (maxH2.size() < k) {
                maxH2.offer(nums[i]);
            } else {
                if (maxH2.peek() < nums[i]) {
                    right[i] = 1;
                } else {
                    maxH2.poll();
                    maxH2.offer(nums[i]);
                }
            }
        }

        // Count the number of elements that are among the k largest from both sides
        int count = 0;
        for (int i = 0; i < n; i++) {
            if (left[i] == 1 && right[i] == 1) {
                count++;
            }
        }

        return count;
    }

    public static void main(String[] args) {
        // Input array and k
        int[] nums = {2, 3, 6, 5, 2, 3};
        int k = 2;

        // Call the function and print the result
        System.out.println(validElements(nums, k));
    }
}

Python3

import heapq

def validElements(nums, k):
    # Initialize two priority queues to keep track of the k largest elements
    max_heap = []
    max_heap2 = []

    n = len(nums)
    left = [0] * n
    right = [0] * n

    # Iterate over the array from left to right
    for i in range(n):
        if len(max_heap) < k:
            heapq.heappush(max_heap, -nums[i])
        else:
            if -max_heap[0] < nums[i]:
                left[i] = 1
            else:
                heapq.heappop(max_heap)
                heapq.heappush(max_heap, -nums[i])

    # Iterate over the array from right to left
    for i in range(n - 1, -1, -1):
        if len(max_heap2) < k:
            heapq.heappush(max_heap2, -nums[i])
        else:
            if -max_heap2[0] < nums[i]:
                right[i] = 1
            else:
                heapq.heappop(max_heap2)
                heapq.heappush(max_heap2, -nums[i])

    # Count the number of elements that are among the k largest from both sides
    count = 0
    for i in range(n):
        if left[i] and right[i]:
            count += 1

    return count

def main():
    # input array and k
    nums = [2, 3, 6, 5, 2, 3]
    k = 2
    
    # Call the function and print the result
    print(validElements(nums, k))

# Call the main function
if __name__ == "__main__":
    main()

C#

using System;
using System.Collections.Generic;

class MainClass {
    public static void Main (string[] args) {
        int[] arr1 = {1,3,6,5,2,1};
        int[] arr2 = {2,2,2};
        
        Console.WriteLine(CountElements(arr1, 2)); // Output: 2
        Console.WriteLine(CountElements(arr2, 3)); // Output: 0
    }
    
    public static int CountElements(int[] arr, int k) {
        int n = arr.Length;
        int[] left = new int[n];
        int[] right = new int[n];
        
        // Priority queue to track k largest elements
        PriorityQueue<int> maxH = new PriorityQueue<int>((a, b) => b.CompareTo(a));
        // Priority queue to track k largest elements from the right side
        PriorityQueue<int> maxH2 = new PriorityQueue<int>((a, b) => b.CompareTo(a));
        
        // Traversing left to right
        for (int i = 0; i < n; i++) {
            if (maxH.Count < k) {
                maxH.Enqueue(arr[i]);
                left[i] = 0;
            } else if (arr[i] > maxH.Peek()) {
                left[i] = 1;
            }
            else {
                maxH.Dequeue();
                maxH.Enqueue(arr[i]);
            }
        }
        
        // Traversing right to left
        for (int i = n - 1; i >= 0; i--) {
            if (maxH2.Count < k) {
                maxH2.Enqueue(arr[i]);
                right[i] = 0;
            } else if (arr[i] > maxH2.Peek()) {
                right[i] = 1;
            }
            else {
                maxH2.Dequeue();
                maxH2.Enqueue(arr[i]);
            }
        }
        
        // Counting elements satisfying conditions on both sides
        int count = 0;
        for (int i = 0; i < n; i++) {
            if (left[i] == 1 && right[i] == 1) {
                count++;
            }
        }
        
        return count;
    }
    
    // Priority Queue implementation
    public class PriorityQueue<T> where T : IComparable<T> {
        private List<T> data;
        private Comparison<T> comparison;

        public PriorityQueue(Comparison<T> comparison) {
            this.data = new List<T>();
            this.comparison = comparison;
        }

        public void Enqueue(T item) {
            data.Add(item);
            int ci = data.Count - 1;
            while (ci > 0) {
                int pi = (ci - 1) / 2;
                if (comparison(data[ci], data[pi]) >= 0) break;
                T tmp = data[ci]; data[ci] = data[pi]; data[pi] = tmp;
                ci = pi;
            }
        }

        public T Dequeue() {
            if (data.Count == 0) throw new InvalidOperationException("Queue is empty");
            int li = data.Count - 1;
            T frontItem = data[0];
            data[0] = data[li];
            data.RemoveAt(li);

            --li;
            int pi = 0;
            while (true) {
                int ci = pi * 2 + 1;
                if (ci > li) break;
                int rc = ci + 1;
                if (rc <= li && comparison(data[rc], data[ci]) < 0) ci = rc;
                if (comparison(data[pi], data[ci]) <= 0) break;
                T tmp = data[pi]; data[pi] = data[ci]; data[ci] = tmp;
                pi = ci;
            }
            return frontItem;
        }

        public T Peek() {
            if (data.Count == 0) throw new InvalidOperationException("Queue is empty");
            return data[0];
        }

        public int Count { get { return data.Count; } }
    }
}

JavaScript

// Define a min heap class
class MinHeap {
    constructor() {
        this.heap = [];
    }

    // Add an element to the heap
    push(val) {
        this.heap.push(val);
        this.bubbleUp(this.heap.length - 1);
    }

    // Remove and return the minimum element from the heap
    pop() {
        if (this.heap.length === 0) return undefined;
        const min = this.heap[0];
        const last = this.heap.pop();
        if (this.heap.length > 0) {
            this.heap[0] = last;
            this.bubbleDown(0);
        }
        return min;
    }

    // Move the element up in the heap
    bubbleUp(index) {
        while (index > 0) {
            const parentIndex = Math.floor((index - 1) / 2);
            if (this.heap[parentIndex] <= this.heap[index]) break;
            [this.heap[parentIndex], this.heap[index]] = [this.heap[index], this.heap[parentIndex]];
            index = parentIndex;
        }
    }

    // Move the element down in the heap
    bubbleDown(index) {
        const length = this.heap.length;
        while (index < length) {
            let minIndex = index;
            const leftIndex = 2 * index + 1;
            const rightIndex = 2 * index + 2;
            if (leftIndex < length && this.heap[leftIndex] < this.heap[minIndex]) {
                minIndex = leftIndex;
            }
            if (rightIndex < length && this.heap[rightIndex] < this.heap[minIndex]) {
                minIndex = rightIndex;
            }
            if (minIndex === index) break;
            [this.heap[minIndex], this.heap[index]] = [this.heap[index], this.heap[minIndex]];
            index = minIndex;
        }
    }

    // Get the size of the heap
    size() {
        return this.heap.length;
    }
}

// Function to find the valid elements
function validElements(nums, k) {
    const n = nums.length;
    const left = new Array(n).fill(false);
    const right = new Array(n).fill(false);

    // Create min heap and push k elements
    const maxHeap = new MinHeap();
    for (let i = 0; i < n; i++) {
        maxHeap.push(nums[i]);
        if (maxHeap.size() > k) maxHeap.pop();
    }

    // Get the kth largest element
    const kthLargest = maxHeap.pop();

    // Mark elements on the left side that are greater than or equal to kth largest
    for (let i = 0; i < n; i++) {
        if (nums[i] >= kthLargest) left[i] = true;
    }

    // Clear the heap and repeat the process for the right side
    maxHeap.heap = [];
    for (let i = n - 1; i >= 0; i--) {
        maxHeap.push(nums[i]);
        if (maxHeap.size() > k) maxHeap.pop();
    }

    // Get the kth largest element from the right side
    const kthLargestRight = maxHeap.pop();

    // Mark elements on the right side that are greater than or equal to kth largest
    for (let i = n - 1; i >= 0; i--) {
        if (nums[i] >= kthLargestRight) right[i] = true;
    }

    // Count the number of elements that are among the k largest from both sides
    let count = 0;
    for (let i = 0; i < n; i++) {
        if (left[i] && right[i]) count++;
    }

    return count;
}

// Input array and k
let nums = [2, 3, 6, 5, 2, 3];
let k = 2;

// Call the function and print the result
console.log(validElements(nums, k));

Output

2

Time Complexity: O(n log (n))
Auxiliary Space: O(k)