 GeeksforGeeks App
Open App Browser
Continue

# Index of kth set bit in a binary array with update queries

Given a binary array arr[] and q queries of following types:

1. k: find the index of the kth set bit i.e. kth 1 in the array.
2. (x, y): Update arr[x] = y where y can either be a 0 or 1.

Examples:

Input: arr[] = {1, 0, 1, 0, 0, 1, 1, 1}, q = 2
k = 4
(x, y) = (5, 1)
Output:
Index of 4th set bit: 6
Array after updation:
1 0 1 0 0 0 1 1

Approach: A simple solution is to run a loop from 0 to n – 1 and find the kth 1. To update a value, simply do arr[i] = x. The first operation takes O(n) time and second operation takes O(1) time.
Another solution is to create another array and store each 1’s index at the ith index in this array. Index of Kth 1 can now be calculated in O(1) time but update operation takes O(n) time now. This works well if the number of query operations are large and very few updates.
But, what if the number of query and updates are equal? We can perform both the operations in O(log n) time using a Segment Tree to do both operations in O(Logn) time.

Representation of Segment tree:

1. Leaf Nodes are the elements of the input array.
2. Each internal node represents sum merging of the leaf nodes.

An array representation of tree is used to represent Segment Trees. For each node at index i, the left child is at index 2*i+1, right child at 2*i+2 and the parent is at (i-1)/2.

Below is the implementation of the above approach:

## C++

 `// C++ implementation of the approach``#include ``using` `namespace` `std;` `// Function to build the tree``void` `buildTree(``int``* tree, ``int``* a, ``int` `s, ``int` `e, ``int` `idx)``{``    ``// s = starting index``    ``// e = ending index``    ``// a = array storing binary string of numbers``    ``// idx = starting index = 1, for recurring the tree``    ``if` `(s == e) {` `        ``// store each element value at the``        ``// leaf node``        ``tree[idx] = a[s];``        ``return``;``    ``}` `    ``int` `mid = (s + e) / 2;` `    ``// Recurring in two sub portions (left, right)``    ``// to build the segment tree` `    ``// Calling for the left sub portion``    ``buildTree(tree, a, s, mid, 2 * idx);` `    ``// Calling for the right sub portion``    ``buildTree(tree, a, mid + 1, e, 2 * idx + 1);` `    ``// Summing up the number of one's``    ``tree[idx] = tree[2 * idx] + tree[2 * idx + 1];` `    ``return``;``}` `// Function to return the index of the query``int` `queryTree(``int``* tree, ``int` `s, ``int` `e, ``int` `k, ``int` `idx)``{``    ``// s = starting index``    ``// e = ending index``    ``// k = searching for kth bit``    ``// idx = starting index = 1, for recurring the tree` `    ``// k > number of 1's in a binary string``    ``if` `(k > tree[idx])``        ``return` `-1;` `    ``// leaf node at which kth 1 is stored``    ``if` `(s == e)``        ``return` `s;``    ``int` `mid = (s + e) / 2;` `    ``// If left sub-tree contains more or equal 1's``    ``// than required kth 1``    ``if` `(tree[2 * idx] >= k)``        ``return` `queryTree(tree, s, mid, k, 2 * idx);` `    ``// If left sub-tree contains less 1's than``    ``// required kth 1 then recur in the right sub-tree``    ``else``        ``return` `queryTree(tree, mid + 1, e, k - tree[2 * idx], 2 * idx + 1);``}` `// Function to perform the update query``void` `updateTree(``int``* tree, ``int` `s, ``int` `e, ``int` `i, ``int` `change, ``int` `idx)``{``    ``// s = starting index``    ``// e = ending index``    ``// i = index at which change is to be done``    ``// change = new changed bit``    ``// idx = starting index = 1, for recurring the tree` `    ``// Out of bounds request``    ``if` `(i < s || i > e) {``        ``cout << ``"error"``;``        ``return``;``    ``}` `    ``// Leaf node of the required index i``    ``if` `(s == e) {` `        ``// Replacing the node value with``        ``// the new changed value``        ``tree[idx] = change;``        ``return``;``    ``}` `    ``int` `mid = (s + e) / 2;` `    ``// If the index i lies in the left sub-tree``    ``if` `(i >= s && i <= mid)``        ``updateTree(tree, s, mid, i, change, 2 * idx);` `    ``// If the index i lies in the right sub-tree``    ``else``        ``updateTree(tree, mid + 1, e, i, change, 2 * idx + 1);` `    ``// Merging both left and right sub-trees``    ``tree[idx] = tree[2 * idx] + tree[2 * idx + 1];``    ``return``;``}` `// Function to perform queries``void` `queries(``int``* tree, ``int``* a, ``int` `q, ``int` `p, ``int` `k, ``int` `change, ``int` `n)``{``    ``int` `s = 0, e = n - 1, idx = 1;``    ``if` `(q == 1) {``        ``// q = 1 update, p = index at which change``        ``// is to be done, change = new bit``        ``a[p] = change;``        ``updateTree(tree, s, e, p, change, idx);``        ``cout << ``"Array after updation:\n"``;``        ``for` `(``int` `i = 0; i < n; i++)``            ``cout << a[i] << ``" "``;``        ``cout << ``"\n"``;``    ``}``    ``else` `{``        ``// q = 0, print kth bit``        ``cout << ``"Index of "` `<< k << ``"th set bit: "``             ``<< queryTree(tree, s, e, k, idx) << ``"\n"``;``    ``}``}` `// Driver code``int` `main()``{``    ``int` `a[] = { 1, 0, 1, 0, 0, 1, 1, 1 };``    ``int` `n = ``sizeof``(a) / ``sizeof``(``int``);` `    ``// Declaring & initializing the tree with``    ``// maximum possible size of the segment tree``    ``// and each value initially as 0``    ``int``* tree = ``new` `int``[4 * n + 1];``    ``for` `(``int` `i = 0; i < 4 * n + 1; ++i) {``        ``tree[i] = 0;``    ``}` `    ``// s and e are the starting and ending``    ``// indices respectively``    ``int` `s = 0, e = n - 1, idx = 1;` `    ``// Build the segment tree``    ``buildTree(tree, a, s, e, idx);` `    ``// Find index of kth set bit``    ``int` `q = 0, p = 0, change = 0, k = 4;``    ``queries(tree, a, q, p, k, change, n);` `    ``// Update query``    ``q = 1, p = 5, change = 0;``    ``queries(tree, a, q, p, k, change, n);` `    ``return` `0;``}`

## Java

 `// Java implementation of the approach``import` `java.io.*;``import` `java.util.*;` `class` `GFG``{` `    ``// Function to build the tree``    ``static` `void` `buildTree(``int``[] tree, ``int``[] a,``                       ``int` `s, ``int` `e, ``int` `idx)``    ``{` `        ``// s = starting index``        ``// e = ending index``        ``// a = array storing binary string of numbers``        ``// idx = starting index = 1, for recurring the tree``        ``if` `(s == e)``        ``{` `            ``// store each element value at the``            ``// leaf node``            ``tree[idx] = a[s];``            ``return``;``        ``}` `        ``int` `mid = (s + e) / ``2``;` `        ``// Recurring in two sub portions (left, right)``        ``// to build the segment tree` `        ``// Calling for the left sub portion``        ``buildTree(tree, a, s, mid, ``2` `* idx);` `        ``// Calling for the right sub portion``        ``buildTree(tree, a, mid + ``1``, e, ``2` `* idx + ``1``);` `        ``// Summing up the number of one's``        ``tree[idx] = tree[``2` `* idx] + tree[``2` `* idx + ``1``];` `        ``return``;``    ``}` `    ``// Function to return the index of the query``    ``static` `int` `queryTree(``int``[] tree, ``int` `s,``                        ``int` `e, ``int` `k, ``int` `idx)``    ``{` `        ``// s = starting index``        ``// e = ending index``        ``// k = searching for kth bit``        ``// idx = starting index = 1, for recurring the tree` `        ``// k > number of 1's in a binary string``        ``if` `(k > tree[idx])``            ``return` `-``1``;` `        ``// leaf node at which kth 1 is stored``        ``if` `(s == e)``            ``return` `s;``        ``int` `mid = (s + e) / ``2``;` `        ``// If left sub-tree contains more or equal 1's``        ``// than required kth 1``        ``if` `(tree[``2` `* idx] >= k)``            ``return` `queryTree(tree, s, mid, k, ``2` `* idx);` `        ``// If left sub-tree contains less 1's than``        ``// required kth 1 then recur in the right sub-tree``        ``else``            ``return` `queryTree(tree, mid + ``1``, e, k - tree[``2` `* idx],``                                                    ``2` `* idx + ``1``);``    ``}` `    ``// Function to perform the update query``    ``static` `void` `updateTree(``int``[] tree, ``int` `s, ``int` `e,``                            ``int` `i, ``int` `change, ``int` `idx)``    ``{` `        ``// s = starting index``        ``// e = ending index``        ``// i = index at which change is to be done``        ``// change = new changed bit``        ``// idx = starting index = 1, for recurring the tree` `        ``// Out of bounds request``        ``if` `(i < s || i > e)``        ``{``            ``System.out.println(``"Error"``);``            ``return``;``        ``}` `        ``// Leaf node of the required index i``        ``if` `(s == e)``        ``{` `            ``// Replacing the node value with``            ``// the new changed value``            ``tree[idx] = change;``            ``return``;``        ``}` `        ``int` `mid = (s + e) / ``2``;` `        ``// If the index i lies in the left sub-tree``        ``if` `(i >= s && i <= mid)``            ``updateTree(tree, s, mid, i, change, ``2` `* idx);` `        ``// If the index i lies in the right sub-tree``        ``else``            ``updateTree(tree, mid + ``1``, e, i, change, ``2` `* idx + ``1``);` `        ``// Merging both left and right sub-trees``        ``tree[idx] = tree[``2` `* idx] + tree[``2` `* idx + ``1``];``        ``return``;``    ``}` `    ``// Function to perform queries``    ``static` `void` `queries(``int``[] tree, ``int``[] a, ``int` `q, ``int` `p,``                                ``int` `k, ``int` `change, ``int` `n)``    ``{``        ``int` `s = ``0``, e = n - ``1``, idx = ``1``;``        ``if` `(q == ``1``)``        ``{` `            ``// q = 1 update, p = index at which change``            ``// is to be done, change = new bit``            ``a[p] = change;``            ``updateTree(tree, s, e, p, change, idx);``            ``System.out.println(``"Array after updation:"``);``            ``for` `(``int` `i = ``0``; i < n; i++)``                ``System.out.print(a[i] + ``" "``);``            ``System.out.println();``        ``}``        ``else``        ``{` `            ``// q = 0, print kth bit``            ``System.out.printf(``"Index of %dth set bit: %d\n"``,``                                ``k, queryTree(tree, s, e, k, idx));``        ``}``    ``}` `    ``// Driver Code``    ``public` `static` `void` `main(String[] args)``    ``{``        ``int``[] a = { ``1``, ``0``, ``1``, ``0``, ``0``, ``1``, ``1``, ``1` `};``        ``int` `n = a.length;` `        ``// Declaring & initializing the tree with``        ``// maximum possible size of the segment tree``        ``// and each value initially as 0``        ``int``[] tree = ``new` `int``[``4` `* n + ``1``];``        ``for` `(``int` `i = ``0``; i < ``4` `* n + ``1``; ++i)``        ``{``            ``tree[i] = ``0``;``        ``}` `        ``// s and e are the starting and ending``        ``// indices respectively``        ``int` `s = ``0``, e = n - ``1``, idx = ``1``;` `        ``// Build the segment tree``        ``buildTree(tree, a, s, e, idx);` `        ``// Find index of kth set bit``        ``int` `q = ``0``, p = ``0``, change = ``0``, k = ``4``;``        ``queries(tree, a, q, p, k, change, n);` `        ``// Update query``        ``q = ``1``;``        ``p = ``5``;``        ``change = ``0``;``        ``queries(tree, a, q, p, k, change, n);``    ``}``}` `// This code is contributed by``// sanjeev2552`

## Python3

 `# Python3 implementation of the approach` `# Function to build the tree``def` `buildTree(tree: ``list``, a: ``list``,``                ``s: ``int``, e: ``int``, idx: ``int``):` `    ``# s = starting index``    ``# e = ending index``    ``# a = array storing binary string of numbers``    ``# idx = starting index = 1, for recurring the tree``    ``if` `s ``=``=` `e:` `        ``# store each element value at the``        ``# leaf node``        ``tree[idx] ``=` `a[s]``        ``return` `    ``mid ``=` `(s ``+` `e) ``/``/` `2` `    ``# Recurring in two sub portions (left, right)``    ``# to build the segment tree` `    ``# Calling for the left sub portion``    ``buildTree(tree, a, s, mid, ``2` `*` `idx)` `    ``# Calling for the right sub portion``    ``buildTree(tree, a, mid ``+` `1``, e, ``2` `*` `idx ``+` `1``)` `    ``# Summing up the number of one's``    ``tree[idx] ``=` `tree[``2` `*` `idx] ``+` `tree[``2` `*` `idx ``+` `1``]` `    ``return` `# Function to return the index of the query``def` `queryTree(tree: ``list``, s: ``int``, e: ``int``,``            ``k: ``int``, idx: ``int``) ``-``> ``int``:` `    ``# s = starting index``    ``# e = ending index``    ``# k = searching for kth bit``    ``# idx = starting index = 1, for recurring the tree` `    ``# k > number of 1's in a binary string``    ``if` `k > tree[idx]:``        ``return` `-``1` `    ``# leaf node at which kth 1 is stored``    ``if` `s ``=``=` `e:``        ``return` `s``    ``mid ``=` `(s ``+` `e) ``/``/` `2` `    ``# If left sub-tree contains more or equal 1's``    ``# than required kth 1``    ``if` `tree[``2` `*` `idx] >``=` `k:``        ``return` `queryTree(tree, s, mid, k, ``2` `*` `idx)` `    ``# If left sub-tree contains less 1's than``    ``# required kth 1 then recur in the right sub-tree``    ``else``:``        ``return` `queryTree(tree, mid ``+` `1``, e,``                    ``k ``-` `tree[``2` `*` `idx], ``2` `*` `idx ``+` `1``)` `# Function to perform the update query``def` `updateTree(tree: ``list``, s: ``int``, e: ``int``,``                ``i: ``int``, change: ``int``, idx: ``int``):` `    ``# s = starting index``    ``# e = ending index``    ``# i = index at which change is to be done``    ``# change = new changed bit``    ``# idx = starting index = 1, for recurring the tree` `    ``# Out of bounds request``    ``if` `i < s ``or` `i > e:``        ``print``(``"error"``)``        ``return` `    ``# Leaf node of the required index i``    ``if` `s ``=``=` `e:` `        ``# Replacing the node value with``        ``# the new changed value``        ``tree[idx] ``=` `change``        ``return` `    ``mid ``=` `(s ``+` `e) ``/``/` `2` `    ``# If the index i lies in the left sub-tree``    ``if` `i >``=` `s ``and` `i <``=` `mid:``        ``updateTree(tree, s, mid, i, change, ``2` `*` `idx)` `    ``# If the index i lies in the right sub-tree``    ``else``:``        ``updateTree(tree, mid ``+` `1``, e, i, change, ``2` `*` `idx ``+` `1``)` `    ``# Merging both left and right sub-trees``    ``tree[idx] ``=` `tree[``2` `*` `idx] ``+` `tree[``2` `*` `idx ``+` `1``]``    ``return` `# Function to perform queries``def` `queries(tree: ``list``, a: ``list``, q: ``int``, p: ``int``,``            ``k: ``int``, change: ``int``, n: ``int``):``    ``s ``=` `0``    ``e ``=` `n ``-` `1``    ``idx ``=` `1``    ``if` `q ``=``=` `1``:` `        ``# q = 1 update, p = index at which change``        ``# is to be done, change = new bit``        ``a[p] ``=` `change``        ``updateTree(tree, s, e, p, change, idx)``        ``print``(``"Array after updation:"``)``        ``for` `i ``in` `range``(n):``            ``print``(a[i], end``=``" "``)``        ``print``()``    ``else``:` `        ``# q = 0, print kth bit``        ``print``(``"Index of %dth set bit: %d"` `%` `(k,``                ``queryTree(tree, s, e, k, idx)))` `# Driver Code``if` `__name__ ``=``=` `"__main__"``:``    ``a ``=` `[``1``, ``0``, ``1``, ``0``, ``0``, ``1``, ``1``, ``1``]``    ``n ``=` `len``(a)` `    ``# Declaring & initializing the tree with``    ``# maximum possible size of the segment tree``    ``# and each value initially as 0``    ``tree ``=` `[``0``] ``*` `(``4` `*` `n ``+` `1``)` `    ``# s and e are the starting and ending``    ``# indices respectively``    ``s ``=` `0``    ``e ``=` `n ``-` `1``    ``idx ``=` `1` `    ``# Build the segment tree``    ``buildTree(tree, a, s, e, idx)` `    ``# Find index of kth set bit``    ``q ``=` `0``    ``p ``=` `0``    ``change ``=` `0``    ``k ``=` `4``    ``queries(tree, a, q, p, k, change, n)` `    ``# Update query``    ``q ``=` `1``    ``p ``=` `5``    ``change ``=` `0``    ``queries(tree, a, q, p, k, change, n)` `# This code is contributed by``# sanjeev2552`

## C#

 `// C# implementation of the approach``using` `System;` `class` `GFG``{` `    ``// Function to build the tree``    ``static` `void` `buildTree(``int``[] tree, ``int``[] a,``                        ``int` `s, ``int` `e, ``int` `idx)``    ``{` `        ``// s = starting index``        ``// e = ending index``        ``// a = array storing binary string of numbers``        ``// idx = starting index = 1, for recurring the tree``        ``if` `(s == e)``        ``{` `            ``// store each element value at the``            ``// leaf node``            ``tree[idx] = a[s];``            ``return``;``        ``}` `        ``int` `mid = (s + e) / 2;` `        ``// Recurring in two sub portions (left, right)``        ``// to build the segment tree` `        ``// Calling for the left sub portion``        ``buildTree(tree, a, s, mid, 2 * idx);` `        ``// Calling for the right sub portion``        ``buildTree(tree, a, mid + 1, e, 2 * idx + 1);` `        ``// Summing up the number of one's``        ``tree[idx] = tree[2 * idx] + tree[2 * idx + 1];` `        ``return``;``    ``}` `    ``// Function to return the index of the query``    ``static` `int` `queryTree(``int``[] tree, ``int` `s,``                        ``int` `e, ``int` `k, ``int` `idx)``    ``{` `        ``// s = starting index``        ``// e = ending index``        ``// k = searching for kth bit``        ``// idx = starting index = 1, for recurring the tree` `        ``// k > number of 1's in a binary string``        ``if` `(k > tree[idx])``            ``return` `-1;` `        ``// leaf node at which kth 1 is stored``        ``if` `(s == e)``            ``return` `s;``        ``int` `mid = (s + e) / 2;` `        ``// If left sub-tree contains more or equal 1's``        ``// than required kth 1``        ``if` `(tree[2 * idx] >= k)``            ``return` `queryTree(tree, s, mid, k, 2 * idx);` `        ``// If left sub-tree contains less 1's than``        ``// required kth 1 then recur in the right sub-tree``        ``else``            ``return` `queryTree(tree, mid + 1,``                            ``e, k - tree[2 * idx],``                            ``2 * idx + 1);``    ``}` `    ``// Function to perform the update query``    ``static` `void` `updateTree(``int``[] tree, ``int` `s, ``int` `e,``                            ``int` `i, ``int` `change, ``int` `idx)``    ``{` `        ``// s = starting index``        ``// e = ending index``        ``// i = index at which change is to be done``        ``// change = new changed bit``        ``// idx = starting index = 1, for recurring the tree` `        ``// Out of bounds request``        ``if` `(i < s || i > e)``        ``{``            ``Console.WriteLine(``"Error"``);``            ``return``;``        ``}` `        ``// Leaf node of the required index i``        ``if` `(s == e)``        ``{` `            ``// Replacing the node value with``            ``// the new changed value``            ``tree[idx] = change;``            ``return``;``        ``}` `        ``int` `mid = (s + e) / 2;` `        ``// If the index i lies in the left sub-tree``        ``if` `(i >= s && i <= mid)``            ``updateTree(tree, s, mid, i,``                        ``change, 2 * idx);` `        ``// If the index i lies in the right sub-tree``        ``else``            ``updateTree(tree, mid + 1, e, i,``                       ``change, 2 * idx + 1);` `        ``// Merging both left and right sub-trees``        ``tree[idx] = tree[2 * idx] +``                    ``tree[2 * idx + 1];``        ``return``;``    ``}` `    ``// Function to perform queries``    ``static` `void` `queries(``int``[] tree, ``int``[] a, ``int` `q, ``int` `p,``                                ``int` `k, ``int` `change, ``int` `n)``    ``{``        ``int` `s = 0, e = n - 1, idx = 1;``        ``if` `(q == 1)``        ``{` `            ``// q = 1 update, p = index at which change``            ``// is to be done, change = new bit``            ``a[p] = change;``            ``updateTree(tree, s, e, p, change, idx);``            ``Console.WriteLine(``"Array after updation:"``);``            ``for` `(``int` `i = 0; i < n; i++)``                ``Console.Write(a[i] + ``" "``);``            ``Console.WriteLine();``        ``}``        ``else``        ``{` `            ``// q = 0, print kth bit``            ``Console.Write(``"Index of {0}th set bit: {1}\n"``,``                            ``k, queryTree(tree, s, e, k, idx));``        ``}``    ``}` `    ``// Driver Code``    ``public` `static` `void` `Main(String[] args)``    ``{``        ``int``[] a = { 1, 0, 1, 0, 0, 1, 1, 1 };``        ``int` `n = a.Length;` `        ``// Declaring & initializing the tree with``        ``// maximum possible size of the segment tree``        ``// and each value initially as 0``        ``int``[] tree = ``new` `int``[4 * n + 1];``        ``for` `(``int` `i = 0; i < 4 * n + 1; ++i)``        ``{``            ``tree[i] = 0;``        ``}` `        ``// s and e are the starting and ending``        ``// indices respectively``        ``int` `s = 0, e = n - 1, idx = 1;` `        ``// Build the segment tree``        ``buildTree(tree, a, s, e, idx);` `        ``// Find index of kth set bit``        ``int` `q = 0, p = 0, change = 0, k = 4;``        ``queries(tree, a, q, p, k, change, n);` `        ``// Update query``        ``q = 1;``        ``p = 5;``        ``change = 0;``        ``queries(tree, a, q, p, k, change, n);``    ``}``}` `// This code is contributed by Rajput-Ji`

## Javascript

 ``

Output:

```Index of 4th set bit: 6
Array after updation:
1 0 1 0 0 0 1 1```

Time Complexity: O(logN) per query and O(NlogN) for building segment tree.
Auxiliary Space: O(N)

My Personal Notes arrow_drop_up