Range Queries for number of Armstrong numbers in an array with updates - GeeksforGeeks

07 Oct, 2021
Given an array arr[] of N integers, the task is to perform the following two queries: 
query(start, end): Print the number of Armstrong numbers in the subarray from start to end
update(i, x): Add x to the array element referenced by array index i, that is: arr[i] = x
Examples: 
Input: arr = { 18, 153, 8, 9, 14, 5} Query 1: query(start = 0, end = 4) Query 2: update(i = 3, x = 11) Query 3: query(start = 0, end = 4) Output: 3 2 Explanation In Query 1, 18 -> 1*1 + 8*8 != 18 153 -> 1*1*1 + 5*5*5 + 3*3*3 = 153 8 -> 8 = 8 9 -> 9 = 9 14 -> 1*1 + 4*4 != 14 the subarray [0...4] has 3 Armstrong numbers viz. {18, 153, 8, 9, 14}In Query 2, the value at index 3 is updated to 11, the array arr now is, { 18, 153, 8, 11, 14, 5}In Query 3, 18 -> 1*1 + 8*8 != 18 153 -> 1*1*1 + 5*5*5 + 3*3*3 = 153 8 -> 8 = 8 9 -> 1*1 + 1*1 != 11 14 -> 1*1 + 4*4 != 14 the subarray [0...4] has 2 Armstrong numbers viz. {18, 153, 8, 11, 14}  
Approach: To handle both point updates and range queries, a segment tree is optimal for this purpose.A positive integer of n digits is called an Armstrong number of order n (order is number of digits) if.  
abcd... = pow(a, n) + pow(b, n) + pow(c, n) + pow(d, n) + .... 
In order to check for Armstrong numbers, the idea is to first count number digits (or find order). Let the number of digits be n. For every digit r in input number x, compute r^n. If the sum of all such values is equal to n, then set it to 1 else to 0.
Building the segment tree:  
The problem is now reduced to the subarray sum using segment tree problem.
Now, we can build the segment tree where a leaf node is represented as either 0 (if it is not an Armstrong number) or 1 (if it is Armstrong number).
The internal nodes of the segment tree equal to the sum of its child nodes, thus a node represent the total Armstrong numbers in the range from L to R with range [L, R] falling under this node and the sub-tree underneath it.
Handling Queries and Point Updates:  
Whenever we receive a query from beginning to end, we can query the segment tree for the sum of nodes in the range from start to end, which in turn represents the number of Armstrong numbers in the range start to end.  
To perform a point update and to update the value at index i to x, we check for the following cases: Let the old value of arri be y and the new value be x. Case 1: If x and y both are Armstrong numbers Count of Armstrong numbers in the subarray does not change so we just update array and do not modify the segment treeCase 2: If x and y both are not Armstrong numbers Count of Armstrong numbers in the subarray does not change so we just update array and do not modify the segment treeCase 3: If y is a Armstrong number but x is not Count of Armstrong numbers in the subarray decreases so we update array and add -1 to every range. The index i which is to be updated is a part of in the segment treeCase 4: If y is not an Armstrong number but x is an Armstrong number Count of Armstrong numbers in the subarray increases so we update array and add 1 to every range. The index i which is to be updated is a part of in the segment tree
Case 1: If x and y both are Armstrong numbers Count of Armstrong numbers in the subarray does not change so we just update array and do not modify the segment treeCase 2: If x and y both are not Armstrong numbers Count of Armstrong numbers in the subarray does not change so we just update array and do not modify the segment treeCase 3: If y is a Armstrong number but x is not Count of Armstrong numbers in the subarray decreases so we update array and add -1 to every range. The index i which is to be updated is a part of in the segment treeCase 4: If y is not an Armstrong number but x is an Armstrong number Count of Armstrong numbers in the subarray increases so we update array and add 1 to every range. The index i which is to be updated is a part of in the segment tree
Case 1: If x and y both are Armstrong numbers Count of Armstrong numbers in the subarray does not change so we just update array and do not modify the segment tree
Case 2: If x and y both are not Armstrong numbers Count of Armstrong numbers in the subarray does not change so we just update array and do not modify the segment tree
Case 3: If y is a Armstrong number but x is not Count of Armstrong numbers in the subarray decreases so we update array and add -1 to every range. The index i which is to be updated is a part of in the segment tree
Case 4: If y is not an Armstrong number but x is an Armstrong number Count of Armstrong numbers in the subarray increases so we update array and add 1 to every range. The index i which is to be updated is a part of in the segment tree
Below is the implementation of the above approach:  
C++
Python3
C#
// C++ program to find the number// of Armstrong numbers in a// subarray and performing updates #include <bits/stdc++.h>using namespace std; #define MAX 1000 // Function that return true// if num is armstrong// else return falsebool isArmstrong(int x){    int n = to_string(x).size();    int sum1 = 0;    int temp = x;    while (temp > 0) {        int digit = temp % 10;        sum1 += pow(digit, n);        temp /= 10;    }    if (sum1 == x)        return true;    return false;} // A utility function to get the middle// index from corner indexes.int getMid(int s, int e){    return s + (e - s) / 2;} // Recursive function to get the number// of Armstrong numbers in a given range/* where    st    --> Pointer to segment tree    index --> Index of current node in the              segment tree. Initially 0 is passed              as root is always at index 0    ss & se  --> Starting and ending indexes of              the segment represented by current              node, i.e., st[index]    qs & qe  --> Starting and ending indexes              of query range      */int queryArmstrongUtil(int* st, int ss,                       int se, int qs,                       int qe, int index){    // If segment of this node is a part    // of given range, then return    // the number of Armstrong numbers    // in the segment    if (qs <= ss && qe >= se)        return st[index];     // If segment of this node    // is outside the given range    if (se < qs || ss > qe)        return 0;     // If a part of this segment    // overlaps with the given range    int mid = getMid(ss, se);    return queryArmstrongUtil(               st, ss, mid, qs,               qe, 2 * index + 1)           + queryArmstrongUtil(                 st, mid + 1, se,                 qs, qe, 2 * index + 2);} // Recursive function to update// the nodes which have the given// index in their range./* where    st, si, ss and se are same as getSumUtil()    i --> index of the element to be updated.          This index is in input array.   diff --> Value to be added to all nodes          which have i in range*/void updateValueUtil(int* st, int ss,                     int se, int i,                     int diff, int si){    // Base Case:    // If the input index lies outside    // the range of this segment    if (i < ss || i > se)        return;     // If the input index is in range    // of this node, then update the value    // of the node and its children    st[si] = st[si] + diff;    if (se != ss) {         int mid = getMid(ss, se);        updateValueUtil(st, ss, mid, i,                        diff, 2 * si + 1);        updateValueUtil(st, mid + 1, se,                        i, diff, 2 * si + 2);    }} // Function to update a value in the// input array and segment tree.// It uses updateValueUtil() to update// the value in segment treevoid updateValue(int arr[], int* st,                 int n, int i,                 int new_val){    // Check for erroneous input index    if (i < 0 || i > n - 1) {        printf("Invalid Input");        return;    }     int diff, oldValue;     oldValue = arr[i];     // Update the value in array    arr[i] = new_val;     // Case 1: Old and new values    // both are Armstrong numbers    if (isArmstrong(oldValue)        && isArmstrong(new_val))        return;     // Case 2: Old and new values    // both not Armstrong numbers    if (!isArmstrong(oldValue)        && !isArmstrong(new_val))        return;     // Case 3: Old value was Armstrong,    // new value is non Armstrong    if (isArmstrong(oldValue) && !isArmstrong(new_val)) {        diff = -1;    }     // Case 4: Old value was non Armstrong,    // new_val is Armstrong    if (!isArmstrong(oldValue)        && !isArmstrong(new_val)) {        diff = 1;    }     // Update the values of    // nodes in segment tree    updateValueUtil(        st, 0, n - 1,        i, diff, 0);} // Return number of Armstrong numbers// in range from index qs (query start)// to qe (query end).// It mainly uses queryArmstrongUtil()void queryArmstrong(int* st, int n,                    int qs, int qe){    int ArmstrongInRange        = queryArmstrongUtil(st, 0, n - 1,                             qs, qe, 0);     cout << "Number of Armstrong numbers "         << "in subarray from "         << qs << " to "         << qe << " = "         << ArmstrongInRange << "\n";} // Recursive function that constructs// Segment Tree for array[ss..se].// si is index of current node// in segment tree stint constructSTUtil(int arr[], int ss,                    int se, int* st,                    int si){    // If there is one element in array,    // check if it is Armstrong number    // then store 1 in the segment tree    // else store 0 and return    if (ss == se) {         // if arr[ss] is Armstrong number        if (isArmstrong(arr[ss]))            st[si] = 1;        else            st[si] = 0;         return st[si];    }     // If there are more than one elements,    // then recur for left and right subtrees    // and store the sum of the    // two values in this node    int mid = getMid(ss, se);    st[si] = constructSTUtil(                 arr, ss, mid, st,                 si * 2 + 1)             + constructSTUtil(                   arr, mid + 1, se, st,                   si * 2 + 2);    return st[si];} // Function to construct a segment// tree from given array.// This function allocates memory// for segment tree and// calls constructSTUtil() to// fill the allocated memoryint* constructST(int arr[], int n){    // Allocate memory for segment tree     // Height of segment tree    int x = (int)(ceil(log2(n)));     // Maximum size of segment tree    int max_size = 2 * (int)pow(2, x) - 1;     int* st = new int[max_size];     // Fill the allocated memory st    constructSTUtil(arr, 0, n - 1, st, 0);     // Return the constructed segment tree    return st;} // Driver Codeint main(){     int arr[] = { 18, 153, 8, 9, 14, 5 };    int n = sizeof(arr) / sizeof(arr[0]);     // Build segment tree from given array    int* st = constructST(arr, n);     // Query 1: Query(start = 0, end = 4)    int start = 0;    int end = 4;    queryArmstrong(st, n, start, end);     // Query 2: Update(i = 3, x = 11),    // i.e Update a[i] to x    int i = 3;    int x = 11;    updateValue(arr, st, n, i, x);     // Print array after update    cout << "Array after update: ";    for (int i = 0; i < n; i++)        cout << arr[i] << ", ";    cout << endl;     // Query 3: Query(start = 0, end = 4)    start = 0;    end = 4;    queryArmstrong(st, n, start, end);     return 0;}
# Python3 program to find the number# of Armstrong numbers in a# subarray and performing updatesimport math MAX = 1000 # Function that return true# if num is armstrong# else return falsedef isArmstrong(x):         n = len(str(x))    sum1 = 0    temp = x         while temp > 0:        digit = temp % 10        sum1 += pow(digit, n)        temp = temp // 10         if sum1 == x:        return True    return False # A utility function to get the middle# index from corner indexes.def getMid(s, e):         return s + (e - s) // 2 # Recursive function to get the number# of Armstrong numbers in a given range# where# st --> Pointer to segment tree# index --> Index of current node in the#             segment tree. Initially 0 is passed#             as root is always at index 0# ss & se --> Starting and ending indexes of#             the segment represented by current#             node, i.e., st[index]# qs & qe --> Starting and ending indexes#             of query rangedef queryArmstrongUtil(st, ss, se, qs, qe, index):         # If segment of this node is a part    # of given range, then return    # the number of Armstrong numbers    # in the segment    if qs <= ss and qe >= se:        return st[index]         # If segment of this node    # is outside the given range    if se < qs or ss > qe:        return 0         # If a part of this segment    # overlaps with the given range    mid = getMid(ss, se)         return (queryArmstrongUtil(st, ss, mid, qs,                               qe, 2 * index + 1) +            queryArmstrongUtil(st, mid + 1, se, qs,                               qe, 2 * index + 2)) # Recursive function to update# the nodes which have the given# index in their range.# where# st, si, ss and se are same as getSumUtil()# i --> index of the element to be updated.#         This index is in input array.# diff --> Value to be added to all nodes#         which have i in rangedef updateValueUtil(st, ss, se, i, diff, si):         # Base Case:    # If the input index lies outside    # the range of this segment    if i < ss or i > se:        return         # If the input index is in range    # of this node, then update the value    # of the node and its children    st[si] = st[si] + diff    if se != ss:        mid = getMid(ss, se)        updateValueUtil(st, ss, mid, i,                        diff, 2 * si + 1)        updateValueUtil(st, mid + 1, se, i,                        diff, 2 * si + 2) # Function to update a value in the# input array and segment tree.# It uses updateValueUtil() to update# the value in segment treedef updateValue(arr, st, n, i, new_val):         # Check for erroneous input index    if i < 0 or i > n - 1:        print('Invalid Input')        return         oldValue = arr[i]         # Update the value in array    arr[i] = new_val         # Case 1: Old and new values    # both are Armstrong numbers    if (isArmstrong(oldValue) and        isArmstrong(new_val)):        return         # Case 2: Old and new values    # both not Armstrong numbers    if (not isArmstrong(oldValue) and        not isArmstrong(new_val)):        return         # Case 3: Old value was Armstrong,    # new value is non Armstrong    if (isArmstrong(oldValue) and (not        isArmstrong(new_val))):        diff = -1         # Case 4: Old value was non Armstrong,    # new_val is Armstrong    if (not isArmstrong(oldValue) and        not isArmstrong(new_val)):        diff = 1         # Update the values of    # nodes in segment tree    updateValueUtil(st, 0, n - 1, i, diff, 0) # Return number of Armstrong numbers# in range from index qs (query start)# to qe (query end).# It mainly uses queryArmstrongUtil()def queryArmstrong(st, n, qs, qe):         ArmstrongInRange = queryArmstrongUtil(st, 0, n - 1,                                          qs, qe, 0)    print("Number of Armstrong numbers in "          "subarray from", qs, "to", qe, "=",           ArmstrongInRange) # Recursive function that constructs# Segment Tree for array[ss..se].# si is index of current node# in segment tree stdef constructSTUtil(arr, ss, se, st, si):         # If there is one element in array,    # check if it is Armstrong number    # then store 1 in the segment tree    # else store 0 and return    if ss == se:                 # If arr[ss] is Armstrong number        if isArmstrong(arr[ss]):            st[si] = 1        else:            st[si] = 0                     return st[si]         # If there are more than one elements,    # then recur for left and right subtrees    # and store the sum of the    # two values in this node    mid = getMid(ss, se)    st[si] = (constructSTUtil(arr, ss, mid,                              st, si * 2 + 1) +              constructSTUtil(arr, mid + 1, se,                              st, si * 2 + 2))                                  return st[si] # Function to construct a segment# tree from given array.# This function allocates memory# for segment tree and# calls constructSTUtil() to# fill the allocated memorydef constructST(arr, n):         # Allocate memory for segment tree     # Height of segment tree    x = int(math.ceil(math.log2(n)))         # Maximum size of segment tree    max_size = 2 * int(pow(2, x)) - 1         st = [-1] * max_size         # Fill the allocated memory st    constructSTUtil(arr, 0, n - 1, st, 0)         # Return the constructed segment tree    return st # Driver codearr = [ 18, 153, 8, 9, 14, 5 ]n = len(arr) # Build segment tree from given arrayst = constructST(arr, n) # Query 1: Query(start = 0, end = 4)start = 0end = 4queryArmstrong(st, n, start, end) # Query 2: Update(i = 3, x = 11),# i.e Update a[i] to xi = 3x = 11updateValue(arr, st, n, i, x) # Print array after updateprint("Array after update:", end = " ")for i in range(n):    print(arr[i], end = ", ")     print() # Query 3: Query(start = 0, end = 4)start = 0end = 4queryArmstrong(st, n, start, end) # This code is contributed by stutipathak31jan
// C# program to find the number// of Armstrong numbers in a// subarray and performing updatesusing System; class GFG{     public int MAX = 1000; // Function that return true// if num is armstrong// else return falsestatic bool isArmstrong(int x){    int n = x.ToString().Length;    int sum1 = 0;    int temp = x;         while (temp > 0)    {        int digit = temp % 10;        sum1 += (int)Math.Pow(digit, n);        temp /= 10;    }         if (sum1 == x)        return true;             return false;} // A utility function to get the middle// index from corner indexes.static int getMid(int s, int e){    return s + (e - s) / 2;} // Recursive function to get the number// of Armstrong numbers in a given range/* where    st    --> Pointer to segment tree    index --> Index of current node in the              segment tree. Initially 0 is passed              as root is always at index 0    ss & se  --> Starting and ending indexes of              the segment represented by current              node, i.e., st[index]    qs & qe  --> Starting and ending indexes              of query range    */static int queryArmstrongUtil(int[] st, int ss, int se,                              int qs, int qe, int index){         // If segment of this node is a part    // of given range, then return    // the number of Armstrong numbers    // in the segment    if (qs <= ss && qe >= se)        return st[index];     // If segment of this node    // is outside the given range    if (se < qs || ss > qe)        return 0;     // If a part of this segment    // overlaps with the given range    int mid = getMid(ss, se);    return queryArmstrongUtil(st, ss, mid, qs, qe,                              2 * index + 1) +           queryArmstrongUtil(st, mid + 1, se, qs, qe,                              2 * index + 2);} // Recursive function to update// the nodes which have the given// index in their range./* where    st, si, ss and se are same as getSumUtil()    i --> index of the element to be updated.          This index is in input array.   diff --> Value to be added to all nodes          which have i in range*/static void updateValueUtil(int[] st, int ss, int se,                            int i, int diff, int si){         // Base Case:    // If the input index lies outside    // the range of this segment    if (i < ss || i > se)        return;     // If the input index is in range    // of this node, then update the value    // of the node and its children    st[si] = st[si] + diff;         if (se != ss)    {        int mid = getMid(ss, se);        updateValueUtil(st, ss, mid, i, diff,                        2 * si + 1);        updateValueUtil(st, mid + 1, se, i, diff,                        2 * si + 2);    }} // Function to update a value in the// input array and segment tree.// It uses updateValueUtil() to update// the value in segment treestatic void updateValue(int[] arr, int[] st, int n,                        int i, int new_val){         // Check for erroneous input index    if (i < 0 || i > n - 1)    {        Console.Write("Invalid Input");        return;    }     int diff = 0, oldValue = 0;     oldValue = arr[i];     // Update the value in array    arr[i] = new_val;     // Case 1: Old and new values    // both are Armstrong numbers    if (isArmstrong(oldValue) &&        isArmstrong(new_val))        return;     // Case 2: Old and new values    // both not Armstrong numbers    if (!isArmstrong(oldValue) &&        !isArmstrong(new_val))        return;     // Case 3: Old value was Armstrong,    // new value is non Armstrong    if (isArmstrong(oldValue) &&        !isArmstrong(new_val))    {        diff = -1;    }     // Case 4: Old value was non Armstrong,    // new_val is Armstrong    if (!isArmstrong(oldValue) &&        !isArmstrong(new_val))    {        diff = 1;    }     // Update the values of    // nodes in segment tree    updateValueUtil(st, 0, n - 1, i, diff, 0);} // Return number of Armstrong numbers// in range from index qs (query start)// to qe (query end).// It mainly uses queryArmstrongUtil()static void queryArmstrong(int[] st, int n, int qs,                           int qe){    int ArmstrongInRange = queryArmstrongUtil(        st, 0, n - 1, qs, qe, 0);     Console.WriteLine("Number of Armstrong numbers " +                      "in subarray from " + qs + " to " +                      qe + " = " + ArmstrongInRange);} // Recursive function that constructs// Segment Tree for array[ss..se].// si is index of current node// in segment tree ststatic int constructSTUtil(int[] arr, int ss, int se,                           int[] st, int si){         // If there is one element in array,    // check if it is Armstrong number    // then store 1 in the segment tree    // else store 0 and return    if (ss == se)    {                 // If arr[ss] is Armstrong number        if (isArmstrong(arr[ss]))            st[si] = 1;        else            st[si] = 0;         return st[si];    }     // If there are more than one elements,    // then recur for left and right subtrees    // and store the sum of the    // two values in this node    int mid = getMid(ss, se);    st[si] = constructSTUtil(arr, ss, mid,                             st, si * 2 + 1) +             constructSTUtil(arr, mid + 1, se,                             st, si * 2 + 2);    return st[si];} // Function to construct a segment// tree from given array.// This function allocates memory// for segment tree and// calls constructSTUtil() to// fill the allocated memorystatic int[] constructST(int[] arr, int n){         // Allocate memory for segment tree     // Height of segment tree    int x = (int)(Math.Ceiling(Math.Log(n, 2)));     // Maximum size of segment tree    int max_size = 2 * (int)Math.Pow(2, x) - 1;     int[] st = new int[max_size];     // Fill the allocated memory st    constructSTUtil(arr, 0, n - 1, st, 0);     // Return the constructed segment tree    return st;} // Driver Codepublic static void Main(string[] args){    int[] arr = { 18, 153, 8, 9, 14, 5 };    int n = arr.Length;     // Build segment tree from given array    int[] st = constructST(arr, n);     // Query 1: Query(start = 0, end = 4)    int start = 0;    int end = 4;    queryArmstrong(st, n, start, end);     // Query 2: Update(i = 3, x = 11),    // i.e Update a[i] to x    int i = 3;    int x = 11;    updateValue(arr, st, n, i, x);     // Print array after update    Console.Write("Array after update: ");    for(int j = 0; j < n; j++)        Console.Write(arr[j] + ", ");             Console.WriteLine();     // Query 3: Query(start = 0, end = 4)    start = 0;    end = 4;    queryArmstrong(st, n, start, end);}} // This code is contributed by ukasp
Number of Armstrong numbers in subarray from 0 to 4 = 3
Array after update: 18, 153, 8, 11, 14, 5, 
Number of Armstrong numbers in subarray from 0 to 4 = 2
 
Time Complexity: The time complexity of each query and update is O(log N) and that of building the segment tree is O(N) 
stutipathak31jan
ukasp
ruhelaa48
array-range-queries
Segment-Tree
subarray
Arrays
Competitive Programming
Divide and Conquer
Write From Home
Arrays
Divide and Conquer
Segment-Tree
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.
Maximum and minimum of an array using minimum number of comparisons
Top 50 Array Coding Problems for Interviews
Stack Data Structure (Introduction and Program)
Introduction to Arrays
Multidimensional Arrays in Java
Competitive Programming - A Complete Guide
Practice for cracking any coding interview
Arrow operator -> in C/C++ with Examples
Prefix Sum Array - Implementation and Applications in Competitive Programming
Fast I/O for Competitive Programming