Open In App

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

Improve
Improve
Like Article
Like
Save
Share
Report

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:

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. 
    1. 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
    2. 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
    3. 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
    4. 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++




// 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 false
bool 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 tree
void 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 st
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 memory
int* 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 Code
int 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;
}


Java




// Java program to find the number
// of Armstrong numbers in a
// subarray and performing updates
import java.util.*;
 
class Main {
 
 
    // Function that return true
    // if num is armstrong
    // else return false
    static boolean isArmstrong(int x) {
        int n = String.valueOf(x).length();
        int sum1 = 0;
        int temp = x;
        while (temp > 0) {
            int digit = temp % 10;
            sum1 += Math.pow(digit, n);
            temp /= 10;
        }
        if (sum1 == x)
            return true;
        return false;
    }
     
    // 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 = (ss + se) / 2;
            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 tree
    static void updateValue(int[] arr, int[] st, int n, int i, int new_val) {
         
        // Check for erroneous input index
        if (i < 0 || i > n - 1) {
            System.out.println("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;
        else
            diff = 1;
         
        // Update the values of
        // nodes in segment tree
        updateValueUtil(st, 0, n - 1, i, diff, 0);
    }
 
 
    // 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 si) {
         
        // 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[si];
         
        // 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 = (ss + se) / 2;
        return queryArmstrongUtil(st, ss, mid, qs, qe, 2 * si + 1) + queryArmstrongUtil(st, mid + 1, se, qs, qe, 2 * si + 2);
    }
     
     
    // 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);
 
        System.out.println("Number of Armstrong numbers in subarray from " + qs + " to " + qe + " = " + ArmstrongInRange);
    }
 
    static int getMid(int s, int e) {
        return s + (e - s) / 2;
    }
     
    // Recursive function that constructs
    // Segment Tree for array[ss..se].
    // si is index of current node
    // in segment tree st
    static 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 memory
    public static int[] constructST(int[] arr, int n) {
         
        // Allocate memory for segment tree
     
        // Height of segment tree
        int x = (int)(Math.ceil(Math.log(n) / Math.log(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;
    }
    public 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);
     
        System.out.print("Array after update: ");
        for (int j = 0; j < n; j++) {
            System.out.print(arr[j] + ", ");
        }
        System.out.println();
         
        // Query 3: Query(start = 0, end = 4)
        start = 0;
        end = 4;
        queryArmstrong(st, n, start, end);
    }
}
// This code is contributed by shivhack999


Python3




# Python3 program to find the number
# of Armstrong numbers in a
# subarray and performing updates
import math
 
MAX = 1000
 
# Function that return true
# if num is armstrong
# else return false
def 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 range
def 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 range
def 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 tree
def 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 st
def 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 memory
def 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 code
arr = [ 18, 153, 8, 9, 14, 5 ]
n = len(arr)
 
# Build segment tree from given array
st = constructST(arr, n)
 
# Query 1: Query(start = 0, end = 4)
start = 0
end = 4
queryArmstrong(st, n, start, end)
 
# Query 2: Update(i = 3, x = 11),
# i.e Update a[i] to x
i = 3
x = 11
updateValue(arr, st, n, i, x)
 
# Print array after update
print("Array after update:", end = " ")
for i in range(n):
    print(arr[i], end = ", ")
     
print()
 
# Query 3: Query(start = 0, end = 4)
start = 0
end = 4
queryArmstrong(st, n, start, end)
 
# This code is contributed by stutipathak31jan


C#




// C# program to find the number
// of Armstrong numbers in a
// subarray and performing updates
using System;
 
class GFG{
     
public int MAX = 1000;
 
// Function that return true
// if num is armstrong
// else return false
static 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 tree
static 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 st
static 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 memory
static 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 Code
public 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


Javascript




// Function that return true
// if num is armstrong
// else return false
function isArmstrong(x){
     
    let n = x.toString().length;
    let sum1 = 0;
    let temp = x;
     
    while (temp > 0) {
        let digit = temp % 10;
        sum1 += Math.pow(digit, n);
        temp = Math.floor(temp / 10);
    }
     
    if (sum1 == x) {
        return true;
    } else {
        return false;
    }
}
 
// A utility function to get the middle
// index from corner indexes.
function getMid(s, e){
     
    return s + Math.floor((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
function 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 && 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
    let 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
function updateValueUtil(st, ss, se, i, diff, 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) {
        let 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 tree
function updateValue(arr, st, n, i, new_val){
     
    // Check for erroneous input index
    if (i < 0 || i > n - 1) {
        console.log("Invalid Input");
        return;
    }
     
    let 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
    let diff = 0;
    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()
function queryArmstrong(st, n, qs, qe){
     
    let ArmstrongInRange = queryArmstrongUtil(st, 0, n - 1,
                              qs, qe, 0);
    console.log("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 st
function 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
    let 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 memory
function constructST(arr, n){
     
    // Allocate memory for segment tree
 
    // Height of segment tree
    let x = Math.ceil(Math.log2(n));
     
    // Maximum size of segment tree
    let max_size = 2 * Math.pow(2, x) - 1;
     
    let st = Array(max_size).fill(-1);
     
    // Fill the allocated memory st
    constructSTUtil(arr, 0, n - 1, st, 0);
     
    // Return the constructed segment tree
    return st;
}
 
// Driver code
let arr = [ 18, 153, 8, 9, 14, 5 ];
let n = arr.length;
 
// Build segment tree from given array
let st = constructST(arr, n);
 
// Query 1: Query(start = 0, end = 4)
let start = 0;
let end = 4;
queryArmstrong(st, n, start, end);
 
// Query 2: Update(i = 3, x = 11),
// i.e Update a[i] to x
let i = 3;
let x = 11;
updateValue(arr, st, n, i, x);
 
// Print array after update
console.log("Array after update: ", end = " ");
for (let i = 0; i < n; i++) {
    console.log(arr[i]);
}
 
console.log();
 
// Query 3: Query(start = 0, end = 4)
start = 0;
end = 4;
queryArmstrong(st, n, start, end);


Output: 

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)
Space Complexity: O(MAX + log x),  where MAX is defined as 1000 and O(log x) for the isArmstrong function.



Last Updated : 25 Apr, 2023
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads