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

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}

Recommended: Please try your approach on {IDE} first, before moving on to the solution.

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; 
} 

Output:

4
2

Time Complexity: O(Q * √N)

My Personal Notes

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

Recommended: Please try your approach on {IDE} first, before moving on to the solution.

C++

Recommended Posts: