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}

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.

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

    filter_none

    edit
    close

    play_arrow

    link
    brightness_4
    code

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

    chevron_right

    
    

    Output:

    4
    2
    

    Time Complexity: O(Q * √N)

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

    competitive-programming-img




    My Personal Notes arrow_drop_up

    Check out this Author's contributed articles.

    If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

    Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.



    Improved By : Akanksha_Rai