Range Queries for count of Armstrong numbers in subarray using MO's algorithm

Range Queries for count of Armstrong numbers in subarray using MO’s algorithm

Last Updated : 08 Nov, 2020

Given an array arr[] consisting of N elements and Q queries represented by L and R denoting a range, the task is to print the number of Armstrong numbers in the subarray [L, R].

Examples:

Input: arr[] = {18, 153, 8, 9, 14, 5}
Query 1: query(L=0, R=5)
Query 2: query(L=3, R=5)
Output: 4
2
Explanation:
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
Query 1: The subarray[0…5] has 4 Armstrong numbers viz. {153, 8, 9, 5}
Query 2: The subarray[3…5] has 2 Armstrong numbers viz. {9, 5}

Approach:
The idea is to use MO’s algorithm to pre-process all queries so that result of one query can be used in the next query.

Sort all queries in a way that queries with L values from 0 to √n – 1 are put together, followed by queries from √n to 2 ×√n – 1, and so on. All queries within a block are sorted in increasing order of R values.
Process all queries one by one and increase the count of Armstrong numbers and store the result in the structure.
- Let ‘count_Armstrong‘ store the count of Armstrong numbers in the previous query.
- Remove extra elements of previous query and add new elements for the current query. For example, if previous query was [0, 8] and the current query is [3, 9], then remove the elements arr[0], arr[1] and arr[2] and add arr[9].
In order to display the results, sort the queries in the order they were provided.

Adding elements

If the current element is a Armstrong number then increment count_Armstrong.

Removing elements

If the current element is a Armstrong number then decrement count_Armstrong.

Below is the implementation of the above approach:

C++

// C++ implementation to count the
// number of Armstrong numbers
// in subarray using MO’s algorithm
 
#include <bits/stdc++.h>
using namespace std;
 
// Variable to represent block size.
// This is made global so compare()
// of sort can use it.
int block;
 
// Structure to represent
// a query range
struct Query {
    int L, R, index;
 
    // Count of Armstrong
    // numbers
    int armstrong;
};
 
// To store the count of
// Armstrong numbers
int count_Armstrong;
 
// Function used to sort all queries so that
// all queries of the same block are arranged
// together and within a block, queries are
// sorted in increasing order of R values.
bool compare(Query x, Query y)
{
    // Different blocks, sort by block.
    if (x.L / block != y.L / block)
        return x.L / block < y.L / block;
 
    // Same block, sort by R value
    return x.R < y.R;
}
 
// Function used to sort all
// queries in order of their
// index value so that results
// of queries can be printed
// in same order as of input
bool compare1(Query x, Query y)
{
    return x.index < y.index;
}
 
// 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;
}
 
// Function to Add elements
// of current range
void add(int currL, int a[])
{
    // If a[currL] is a Armstrong number
    // then increment count_Armstrong
    if (isArmstrong(a[currL]))
        count_Armstrong++;
}
 
// Function to remove elements
// of previous range
void remove(int currR, int a[])
{
    // If a[currL] is a Armstrong number
    // then decrement count_Armstrong
    if (isArmstrong(a[currR]))
        count_Armstrong--;
}
 
// Function to generate the result of queries
void queryResults(int a[], int n, Query q[],
                  int m)
{
 
    // Initialize count_Armstrong to 0
    count_Armstrong = 0;
 
    // Find block size
    block = (int)sqrt(n);
 
    // Sort all queries so that queries of
    // same blocks are arranged together.
    sort(q, q + m, compare);
 
    // Initialize current L, current R and
    // current result
    int currL = 0, currR = 0;
 
    for (int i = 0; i < m; i++) {
        // L and R values of current range
        int L = q[i].L, R = q[i].R;
 
        // Add Elements of current range
        while (currR <= R) {
            add(currR, a);
            currR++;
        }
        while (currL > L) {
            add(currL - 1, a);
            currL--;
        }
 
        // Remove element of previous range
        while (currR > R + 1)
 
        {
            remove(currR - 1, a);
            currR--;
        }
        while (currL < L) {
            remove(currL, a);
            currL++;
        }
 
        q[i].armstrong = count_Armstrong;
    }
}
 
// Function to display the results of
// queries in their initial order
void printResults(Query q[], int m)
{
    sort(q, q + m, compare1);
    for (int i = 0; i < m; i++) {
        cout << q[i].armstrong << endl;
    }
}
 
// Driver Code
int main()
{
 
    int arr[] = { 18, 153, 8, 9, 14, 5 };
    int n = sizeof(arr) / sizeof(arr[0]);
 
    Query q[] = { { 0, 5, 0, 0 },
                  { 3, 5, 1, 0 } };
 
    int m = sizeof(q) / sizeof(q[0]);
 
    queryResults(arr, n, q, m);
 
    printResults(q, m);
 
    return 0;
}

Java

// Java implementation to count the
// number of Armstrong numbers
// in subarray using MO’s algorithm
import java.io.*;
import java.util.*;
 
// Class to represent
// a query range
class Query
{
    int L, R, index;
 
    // Count of Armstrong
    // numbers
    int armstrong;
 
    Query(int L, int R, int index,
          int armstrong)
    {
        this.L = L;
        this.R = R;
        this.index = index;
        this.armstrong = armstrong;
    }
}
 
class GFG{
     
// Variable to represent block size.
static int block;
 
// To store the count of
// Armstrong numbers
static int count_Armstrong;
 
// 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;
}
 
// Function to Add elements
// of current range
static void add(int currL, int a[])
{
     
    // If a[currL] is a Armstrong number
    // then increment count_Armstrong
    if (isArmstrong(a[currL]))
        count_Armstrong++;
}
 
// Function to remove elements
// of previous range
static void remove(int currR, int a[])
{
     
    // If a[currL] is a Armstrong number
    // then decrement count_Armstrong
    if (isArmstrong(a[currR]))
        count_Armstrong--;
}
 
// Function to generate the result of queries
static void queryResults(int a[], int n, Query q[],
                         int m)
{
     
    // Initialize count_Armstrong to 0
    count_Armstrong = 0;
 
    // Find block size
    block = (int)(Math.sqrt(n));
 
    // sort all queries so that
    // all queries of the same block are arranged
    // together and within a block, queries are
    // sorted in increasing order of R values.
    Arrays.sort(q, (Query x, Query y) ->
    {
         
        // Different blocks, sort by block.
        if (x.L / block != y.L / block)
            return x.L / block - y.L / block;
             
        // Same block, sort by R value
        return x.R - y.R;
    });
 
    // Initialize current L, current R and
    // current result
    int currL = 0, currR = 0;
 
    for(int i = 0; i < m; i++)
    {
         
        // L and R values of current range
        int L = q[i].L, R = q[i].R;
 
        // Add Elements of current range
        while (currR <= R)
        {
            add(currR, a);
            currR++;
        }
        while (currL > L)
        {
            add(currL - 1, a);
            currL--;
        }
 
        // Remove element of previous range
        while (currR > R + 1)
        {
            remove(currR - 1, a);
            currR--;
        }
        while (currL < L)
        {
            remove(currL, a);
            currL++;
        }
 
        q[i].armstrong = count_Armstrong;
    }
}
 
// Function to display the results of
// queries in their initial order
static void printResults(Query q[], int m)
{
    Arrays.sort(q, (Query x, Query y) ->
     
                // sort all queries
                // in order of their
                // index value so that results
                // of queries can be printed
                // in same order as of input);
                x.index - y.index);
 
    for(int i = 0; i < m; i++)
    {
        System.out.println(q[i].armstrong);
    }
}
 
// Driver Code
public static void main(String[] args)
{
    int arr[] = { 18, 153, 8, 9, 14, 5 };
    int n = arr.length;
 
    Query q[] = new Query[2];
    q[0] = new Query(0, 5, 0, 0);
    q[1] = new Query(3, 5, 1, 0);
 
    int m = q.length;
 
    queryResults(arr, n, q, m);
 
    printResults(q, m);
}
}
 
// This code is contributed by jithin

Output:

4
2

Time Complexity: O(Q * √N)
Related article: Range Queries for number of Armstrong numbers in an array with updates

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the DSA Self Paced Course at a student-friendly price and become industry ready. To complete your preparation from learning a language to DS Algo and many more, please refer Complete Interview Preparation Course.

In case you wish to attend live classes with industry experts, please refer DSA Live Classes

My Personal Notes

Related Articles

C++

Java