Open In App

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

Given an array arr[] of N integers, the task is to perform the following two queries: 

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:  

Handling Queries and Point Updates:  

Below is the implementation of the above approach:  




// 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 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 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# 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




// 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.


Article Tags :