2519. Count the Number of K-Big Indices


Problem Description

The problem presents an integer array nums and a positive integer k. We have to determine the number of indices in the array that satisfy a certain "k-big" property. An index i is considered "k-big" if there are at least k elements that are less than nums[i] to the left of i and also at least k elements that are less than nums[i] to the right of i. The task is to calculate the total number of such "k-big" indices.

Intuition

The key to solving this problem efficiently lies in being able to quickly determine how many elements are smaller than a given element to its left and right. A naive approach might try to compare each element to all others, resulting in an impractical time complexity.

To optimize this, we can use a Binary Indexed Tree (BIT), also known as a Fenwick tree. This data structure allows us to efficiently compute prefix sums and update values, which is exactly what we need to keep track of the count of elements smaller than the current one.

The solution approach involves two Binary Indexed Trees:

  • tree1 tracks the number of smaller elements to the left of each index.
  • tree2 tracks the number of smaller elements to the right of each index.

We initially build tree2 with the counts of all elements. As we traverse the elements in nums, we update tree2 to remove the count of the current element, essentially keeping the count of elements to the right. We then check if the current index satisfies the "k-big" condition using both tree1 and tree2. If so, we increment our answer. Simultaneously, we update tree1 to reflect this element for the counts of the subsequent indices.

This way, as we iterate over the array, we accumulate the "k-big" indices by efficiently querying and updating the prefix sums using the Binary Indexed Trees.

Learn more about Segment Tree, Binary Search, Divide and Conquer and Merge Sort patterns.

Solution Approach

The solution applies a Binary Indexed Tree to manage the computation of the counts of elements smaller than a certain value to the left and right of a given index efficiently.

Here are the steps of the implementation:

Binary Indexed Tree (BIT) Structure:

  • A class BinaryIndexedTree is defined to encapsulate the BIT operations such as update and query. Each instance of this class maintains an internal array c, where n+1 is the size of the given nums array.

Update Operation:

  • In the update method of BIT, we add a value delta to a specific index x and then update all subsequent elements in the array c that are responsible for that index. This is done by iterating through the array using the bitwise operation x += x & -x.

Query Operation:

  • The query method calculates the prefix sum up to a given index x. It iterates through the array c while decrementing the index using the bit manipulation x -= x & -x until x is zero. The result is the sum of elements for the range [1, x].

Solution Class and Method:

  • In the Solution class, the method kBigIndices implements the main logic.
  • We initialize two instances of the BinaryIndexedTree: tree1 for counting elements to the left and tree2 for counting elements to the right.

Building tree2:

  • The full tree2 is built with the input array by incrementing the count for each value found in nums.

Traverse the Array:

  • As we iterate over the elements in nums, for each element v:
    • We first decrement the count in tree2 for the current element since we want to exclude it from the right count.
    • Then we check if the element v is "k-big" using both tree1 and tree2:
      • tree1.query(v - 1) checks for at least k elements less than v to the left.
      • tree2.query(v - 1) checks for at least k elements less than v to the right.
    • If both conditions are true, we increment our answer (ans).
    • We update tree1 to include the current element in the left count for future iterations by using tree1.update(v, 1).

Following these steps ensures an efficient computation of the "k-big" indices, with time complexity that allows the problem to be solved in practical scenarios, even for larger input arrays.

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 we have the following integer array nums and integer k:

nums = [3, 5, 1, 2, 4]
k = 2

With this input, we want to find the number of indices where the "k-big" property is satisfied, meaning there are at least 2 elements smaller to both the left and right of that index. Here's how the solution works step by step:

Step 1: Initialize two Binary Indexed Trees (BITs)

  • We will have tree1 to track counts to the left and tree2 to track counts to the right.

Step 2: Building tree2

  • Before we start checking each element, build tree2 with the count of all elements in nums. After populating tree2, it would look something like this (assuming 1-based indexing for the tree and also assuming nums elements as their indices):
tree2 based on values: [0, 1, 1, 1, 1, 1]

Step 3: Traverse the Array and Perform Operations

  • Now, we iterate through each element in nums.

For index 0 (nums[0] = 3):

  • Update tree2 by decrementing the count of the value 3.
  • Check if there are at least k (2) elements less than 3 to the left using tree1.query(2) (which would return 0 right now).
  • Check if there are at least k (2) elements less than 3 to the right using tree2.query(2) (which returns 2, since the values 1 and 2 are less than 3).
  • Condition for "k-big" is not satisfied because we don't have enough elements on the left.
  • Update tree1 with the current value (3) by using tree1.update(3, 1).

For index 1 (nums[1] = 5):

  • Update tree2 by decrementing the count for the value 5.
  • tree1.query(4) checks elements less than 5 to the left, returns 1 (only 3 is less than 5).
  • tree2.query(4) checks elements less than 5 to the right, returns 2 (since 1 and 2 are less than 5).
  • "k-big" condition is not satisfied for the same reason as above.
  • Update tree1 with the current value (5) by using tree1.update(5, 1).

For index 2 (nums[2] = 1):

  • Same procedures as above, but there will be no updates since it's clear there are no elements less than 1 to the left or right.

For index 3 (nums[3] = 2):

  • Update tree2, decrementing the count of value 2.
  • Use tree1.query(1) to check the left side which returns 1, not meeting the "k-big".
  • Use tree2.query(1) to check the right side which returns 1 as well.
  • "k-big" condition is not met.
  • Update tree1 with the current value (2).

For index 4 (nums[4] = 4):

  • Update tree2, decrementing the count of value 4.
  • Use tree1.query(3) to find if there are at least 2 numbers that are < 4 on the left; it returns 2 (since 3 and 2 are less than 4).
  • Use tree2.query(3) to find if there are at least 2 numbers that are < 4 on the right; it returns 0 (since we are at the end, no elements to the right).
  • "k-big" condition is not met.

After we have finished iterating through the elements in nums, we come to realize that none of the indices satisfy the "k-big" property for k = 2. Therefore, our answer would be 0.

This example walked through each element and showed how we check the "k-big" property using the two Binary Indexed Trees. In a larger input array, the efficiency of the BIT operations for update and query becomes crucial, making the solution scalable.

Solution Implementation

1class BinaryIndexedTree:
2    def __init__(self, size):
3        # Initialize the Binary Indexed Tree with size 'size'
4        self.size = size
5        self.tree = [0] * (size + 1)
6
7    def update(self, index, delta):
8        # Update the tree with 'delta' at the given 'index'
9        while index <= self.size:
10            self.tree[index] += delta
11            index += index & -index
12
13    def query(self, index):
14        # Query the prefix sum up to the given 'index'
15        total_sum = 0
16        while index:
17            total_sum += self.tree[index]
18            index -= index & -index
19        return total_sum
20
21
22class Solution:
23    def kBigIndices(self, nums, k):
24        # Find the number of indices for which there are at least 'k' numbers smaller before and after the current number.
25        n = len(nums)
26      
27        # Initialize two Binary Indexed Trees
28        tree_before = BinaryIndexedTree(n)
29        tree_after = BinaryIndexedTree(n)
30      
31        # Initially, consider all numbers in the 'tree_after'
32        for number in nums:
33            tree_after.update(number, 1)
34
35        count = 0  # This will hold the number of valid indices
36
37        # Now, process each number in the input 'nums' list
38        for number in nums:
39            tree_after.update(number, -1)  # Exclude the current number from 'tree_after'
40            # Query if there are at least 'k' numbers smaller before the current number in the 'tree_before'
41            # and 'k' numbers smaller after the current number in the 'tree_after'
42            if tree_before.query(number - 1) >= k and tree_after.query(number - 1) >= k:
43                count += 1
44
45            # Include the current number in the 'tree_before'
46            tree_before.update(number, 1)
47
48        return count
49
50# Example of usage:
51# solution = Solution()
52# result = solution.kBigIndices([3,5,1,2,4], 2)
53# print(result)  # This will print the number of indices with at least 2 numbers smaller before and after.
54
1// Implementation of a Binary Indexed Tree (also known as Fenwick Tree)
2class BinaryIndexedTree {
3    private final int size; // Total number of elements
4    private final int[] tree; // The tree represented as an array
5
6    public BinaryIndexedTree(int size) {
7        this.size = size;
8        this.tree = new int[size + 1]; // Indexing starts at 1, so we need an extra space
9    }
10
11    // Updates the tree with the value 'delta' at the given index 'index'
12    public void update(int index, int delta) {
13        while (index <= size) {
14            tree[index] += delta; // Increment the value at the index
15            index += index & -index; // Move index to parent node
16        }
17    }
18
19    // Queries the cumulative frequency up to and including the given index 'index'
20    public int query(int index) {
21        int sum = 0;
22        while (index > 0) {
23            sum += tree[index]; // Add the value at the index to the sum
24            index -= index & -index; // Move to parent node
25        }
26        return sum; // Return the cumulative frequency
27    }
28}
29
30class Solution {
31    public int kBigIndices(int[] nums, int k) {
32        int n = nums.length;
33        // Create two Binary Indexed Trees, 
34        // treePrefix will count the frequencies of each number in the prefix of nums
35        BinaryIndexedTree treePrefix = new BinaryIndexedTree(n);
36        // treeSuffix will count the frequencies of each number in the suffix of nums
37        BinaryIndexedTree treeSuffix = new BinaryIndexedTree(n);
38      
39        // Initialize treeSuffix with the frequencies of each number in nums
40        for (int value : nums) {
41            treeSuffix.update(value, 1);
42        }
43      
44        int count = 0; // Initialize the count of k-big indices
45      
46        // Iterate over the array 'nums' and use the trees to find k-big indices
47        for (int value : nums) {
48            // Decrease frequency of the current number in treeSuffix because we are moving it to treePrefix
49            treeSuffix.update(value, -1);
50          
51            // Check if the current number is a k-big index
52            if (treePrefix.query(value - 1) >= k && treeSuffix.query(value - 1) >= k) {
53                count++; // Increment the count if both the number of smaller numbers is at least 'k' on both sides
54            }
55          
56            // Move the current number from suffix to prefix
57            treePrefix.update(value, 1);
58        }
59      
60        return count; // Return the total count of k-big indices
61    }
62}
63
1class BinaryIndexedTree {
2public:
3    // Constructor to initialize the BinaryIndexedTree with a given size
4    BinaryIndexedTree(int size)
5        : size_(size)
6        , tree_(size + 1, 0) {} // Initializes the tree with zeros
7
8    // Function to update the tree by adding 'delta' at the given index 'idx'
9    void update(int idx, int delta) {
10        while (idx <= size_) {
11            tree_[idx] += delta; // Add 'delta' to the current index
12            idx += (idx & -idx); // Move to the parent in the tree
13        }
14    }
15
16    // Function to query the cumulative frequency up to the given index 'idx'
17    int query(int idx) const {
18        int total = 0;
19        while (idx > 0) {
20            total += tree_[idx]; // Accumulate the frequencies
21            idx -= (idx & -idx); // Move to the next element to add
22        }
23        return total; // Return the cumulative frequency
24    }
25
26private:
27    int size_;                // Size of the original array
28    std::vector<int> tree_;   // Binary Indexed Tree storage
29};
30
31class Solution {
32public:
33    // Function to return the number of indices such that there are at least 'k' elements
34    // less than or equal to nums[i] both before and after the current index in 'nums'
35    int kBigIndices(std::vector<int>& nums, int k) {
36        int n = nums.size();
37        // Create two Binary Indexed Trees to keep track of frequencies
38        BinaryIndexedTree beforeTree(n);
39        BinaryIndexedTree afterTree(n);
40      
41        // Initialize the 'afterTree' with the frequencies of all elements
42        for (const int& value : nums) {
43            afterTree.update(value, 1);
44        }
45
46        int count = 0; // Variable to store the answer
47      
48        // Iterate through 'nums' and use 'beforeTree' and 'afterTree' to find the answer
49        for (const int& value : nums) {
50            afterTree.update(value, -1); // Update 'afterTree' as we process each element
51            // Check if there are at least 'k' elements that are less than or equal to 'value' before and after the index
52            if (beforeTree.query(value - 1) >= k && afterTree.query(value - 1) >= k) {
53                count++; // Increment the count if the condition is met
54            }
55            beforeTree.update(value, 1); // Update 'beforeTree' as we move to next element
56        }
57      
58        return count; // Return the final count
59    }
60};
61
1// A representation of the Binary Indexed Tree (BIT) as an array
2let tree: number[];
3
4// Initialize the Binary Indexed Tree with a given size
5function initializeBIT(size: number) {
6  tree = new Array(size + 1).fill(0);
7}
8
9// Update the BIT by adding 'delta' at the given index 'idx'
10function updateBIT(idx: number, delta: number, size: number) {
11  while (idx <= size) {
12    tree[idx] += delta; // Add 'delta' to the current index
13    idx += (idx & -idx); // Move to the parent in the tree
14  }
15}
16
17// Query the cumulative frequency up to the given index 'idx'
18function queryBIT(idx: number): number {
19  let total = 0;
20  while (idx > 0) {
21    total += tree[idx]; // Accumulate the frequencies
22    idx -= (idx & -idx); // Move to the next element to add
23  }
24  return total; // Return the cumulative frequency
25}
26
27// Variable to store the size of the original array
28let size: number;
29
30// Function to return the number of indices such that there are at least 'k' elements
31// less than or equal to nums[i] both before and after the current index in 'nums'
32function kBigIndices(nums: number[], k: number): number {
33  size = nums.length;
34  // Initialize before and after BIT with zeros
35  initializeBIT(size);
36  let beforeTree = [...tree];
37  initializeBIT(size);
38  let afterTree = [...tree];
39
40  // Initialize the 'afterTree' with the frequencies of all elements
41  nums.forEach((value) => {
42    updateBIT(value, 1, size);
43  });
44  afterTree = [...tree];
45
46  let count = 0; // Variable to store the answer
47
48  // Iterate through 'nums' and use 'beforeTree' and 'afterTree' to find the answer
49  for (let i = 0; i < nums.length; i++) {
50    const value = nums[i];
51    tree = [...afterTree];
52    updateBIT(value, -1, size);
53    afterTree = [...tree]; // Update 'afterTree' as we process each element
54
55    // Check if there are at least 'k' elements that are less than or equal to 'value' before and after the index
56    tree = beforeTree;
57    const before = queryBIT(value - 1);
58    tree = afterTree;
59    const after = queryBIT(value - 1);
60    if (before >= k && after >= k) {
61      count++; // Increment the count if the condition is met
62    }
63
64    tree = beforeTree;
65    updateBIT(value, 1, size);
66    beforeTree = [...tree]; // Update 'beforeTree' as we move to the next element
67  }
68
69  return count; // Return the final count
70}
71

Time and Space Complexity

The provided code uses two Binary Indexed Trees (BITs) to store and update the frequency of elements in nums. The method update will be called n times for each tree, and each call takes O(log n) time as it iterates through the tree updating nodes. The query method is also called n times for each tree within the for loop, and it likewise takes O(log n) time to aggregate values from the tree.

Given that we have two BITs and each operation (update and query) is O(log n), the overall time complexity is O(n log n), where n is the number of elements in nums.

For space complexity, each tree has an array of size n + 1 to store cumulative frequencies, which results in O(n) space complexity. Since there are two such arrays, the total space used by the data structures is O(2n) which simplifies to O(n).

The code will trigger these operations:

  • update(x, delta) is called 2n times, resulting in O(2n log n) operations.
  • query(x) is called 2n times, also resulting in O(2n log n) operations.

When combined and simplified, the computational complexity remains O(n log n) for time and O(n) for space. The reference answer aligns with this analysis.

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's the output of running the following function using the following tree as input?

1def serialize(root):
2    res = []
3    def dfs(root):
4        if not root:
5            res.append('x')
6            return
7        res.append(root.val)
8        dfs(root.left)
9        dfs(root.right)
10    dfs(root)
11    return ' '.join(res)
12
1import java.util.StringJoiner;
2
3public static String serialize(Node root) {
4    StringJoiner res = new StringJoiner(" ");
5    serializeDFS(root, res);
6    return res.toString();
7}
8
9private static void serializeDFS(Node root, StringJoiner result) {
10    if (root == null) {
11        result.add("x");
12        return;
13    }
14    result.add(Integer.toString(root.val));
15    serializeDFS(root.left, result);
16    serializeDFS(root.right, result);
17}
18
1function serialize(root) {
2    let res = [];
3    serialize_dfs(root, res);
4    return res.join(" ");
5}
6
7function serialize_dfs(root, res) {
8    if (!root) {
9        res.push("x");
10        return;
11    }
12    res.push(root.val);
13    serialize_dfs(root.left, res);
14    serialize_dfs(root.right, res);
15}
16

Recommended Readings

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