Counting Distinct Arrays by Removal and Concatenation of Elements

Given an array arr[] of length N, the task is to create an array res[] of length N where each element res[i] represents the count of distinct arrays obtained by applying the below operation on prefix arr[1, i] for (1 <= i <= N):

• Select three distinct indices let say i₁, i₂ and i₃ such that (1 <= i₁ < i₂ < i₃ <= i) and arr[i1] > arr[i2] < arr[i3]. Remove arr[i₂] and concatenate the remaining elements without changing the order.

Note: Number of distinct arrays can be very large. Therefore, output it in modulo 998244353.

Examples:

Input: N = 4, arr[] = {4, 1, 3, 2}
Output: res[] = {1, 1, 2, 2}
Explanation: There are four possible prefixes are there for arr[], which are {4}, {4, 1}, {4, 1, 3} and {4, 1, 3, 2} respectively. Let us apply the given operation on each prefix then:

• 1st Prefix: arr[1, 1] = {4}
  • As we need at least three indices to perform operations. Therefore, we can't perform operation on this prefix. Distinct number of arrays for prefix arr[1, 1] is 1. Which is {4} itself.
• 2nd Prefix: arr[1, 2] = {4, 1}
  • As we need at least three indices to perform operations. Therefore, we can't perform operation on this prefix. Distinct number of arrays for prefix arr[1, 2] is 1. Which is {4, 1} itself.
• 3rd Prefix: arr[1, 3] = {4, 1, 3}
  • Select i₁, i₂ and i₃ as 1, 2 and 3 respectively. Then the condition of (1 <= i₁ < i₂ < i₃ <= i) and arr[i₁] > arr[i₂] <a rr[i₃] holds true for arr[1, 3] as (4 > 1 < 3). Therefore, remove arr[i₂] = 1 and concatenate the remaining array, which will be {4, 3}. As no further operation can be applied on {4, 3}. Therefore, the distinct arrays that can be obtained by prefix arr[1, 3] are 2, {4, 1, 3} and {4, 3} respectively.
• 4th Prefix: arr[1, 4] = {4, 1, 3, 2}
  • Select i₁, i₂ and i₃ as 1, 2 and 3 respectively. Then the condition of (1 <= i₁ < i₂ < i₃ <= i) and arr[i₁] > arr[i₂] < arr[i₃] holds true for arr[1, 4] as (4 > 1 < 3). Therefore, remove arr[i₂] = 1 and concatenate the remaining array, which will be {4, 3, 2}. As no further operation can be applied on {4, 3, 2}. Therefore, the distinct arrays that can be obtained by prefix arr[1, 4] are 2, {4, 1, 3, 2} and {4, 3, 2} respectively.

Total number of distinct arrays by each prefix are 1, 1, 2 and 2 respectively. Which stored in res[] as {1, 1, 2, 2}.

Input: N = 6, arr[] = {1, 2, 3, 4, 5, 6}
Output: res[] = {1, 1, 1, 1, 1, 1}
Explanation: It can be verified that for each prefix there will be only one distinct array, which will be the prefix itself.

Approach: Implement the idea below to solve the problem:

The problem is observation based. It must be noted that an element in the array can only be removed if it has both a previous and a next greater element. So, for each prefix of the array, we count the number of elements that meet this condition and calculate the number of distinct arrays that can be generated as 2 to the power of this count. This is because each of these elements can either be included or excluded from the array, giving us 2 choices for each element. The total number of distinct arrays is then the product of these choices for all elements.

So the approach boils down to finding the Next Greater and Previous Greater element for each index of arr[]. After that just keep finding number of elements let say X which have both next greater and previous greater within the current range. For each prefix, the result is simple 2^X where X is number of such elements.

Steps-by-step approach:

• Initialize an array let say have[] of size N to store the count of elements which have both a next greater and a previous greater element within the current range. Along with this initialize a counter variable let say Count.
• Declare two arrays let say Prev[] and Next[] and initialize both of them with -1s.
• Then calculate previous greater and next greater element for each element of arr[] in Prev[] and Next[] respectively.
• Run a loop for i = 0 to i < N and follow below mentioned steps:
  • If (Prev[i] != -1 && Next[i] != -1)
    • Increment Have[Next[i]].
  • Count += Have[i]
  • Output Power(2, Count).

Below is the implementation of the above approach:

#include <iostream>
#include <vector>
#include <stack>
#include <algorithm>

using namespace std;

// Mod value
const int mod = 998244353;

// Function declarations
void Find_Distinct_Arrays(int N, vector<int>& arr);
vector<int> PrevGreater(const vector<int>& arr);
vector<int> NextGreater(const vector<int>& arr);
int binpow(int x, int y);

// Driver Function
int main()
{
    // Inputs
    int N = 4;
    vector<int> arr = { 3, 1, 4, 2 };

    // Function call
    Find_Distinct_Arrays(N, arr);

    return 0;
}

// Function to output array B[]
void Find_Distinct_Arrays(int N, vector<int>& arr)
{
    // Array to store the count of elements which have
    // both next greater and prev greater within the
    // current range
    vector<int> have(N, 0);
    int cnt = 0;

    // Arrays to store previous and next greater
    // elements respectively
    vector<int> prev = PrevGreater(arr);
    vector<int> next = NextGreater(arr);

    for (int i = 0; i < N; i++) {
        // If the current element has both
        // a previous and a next greater element
        if (prev[i] != -1 && next[i] != -1) {
            // Increment the count of elements which
            // have both next greater and prev greater within the
            // current range
            // The count is stored at the index of the
            // next greater element
            have[next[i]]++;
        }
        // Add the count of such elements up to the
        // current index to the running total
        cnt += have[i];
        // Print 2 to the power of the running total,
        // modulo 998244353
        // This is the number of distinct arrays that
        // can be generated through the operations for
        // the current prefix of the permutation
        cout << binpow(2, cnt) << " ";
    }
}

// Function to calculate prev greater element using stack
vector<int> PrevGreater(const vector<int>& arr)
{
    // length of arr[]
    int N = arr.size();

    // Stack to find previous and next greater element
    stack<int> s;

    // Array to store the previous greater element
    vector<int> prev(N, -1);

    // Find previous greater element for each element
    for (int i = 0; i < N; i++) {
        while (!s.empty() && arr[s.top()] < arr[i]) {
            s.pop();
        }
        if (!s.empty()) {
            prev[i] = s.top();
        }
        s.push(i);
    }

    // returning array
    return prev;
}

// Function to calculate next greater element using stack
vector<int> NextGreater(const vector<int>& arr)
{
    // length of arr[]
    int N = arr.size();

    // Stack to find previous and next greater element
    stack<int> s;

    // Array to store next greater element
    vector<int> next(N, -1);

    // Find next greater element for each element
    for (int i = N - 1; i >= 0; i--) {
        while (!s.empty() && arr[s.top()] < arr[i]) {
            s.pop();
        }
        if (!s.empty()) {
            next[i] = s.top();
        }
        s.push(i);
    }

    // returning array
    return next;
}

// Function to calculate power in log(n) time
int binpow(int x, int y)
{
    int res = 1;
    while (y > 0) {
        if (y % 2 == 1)
            res = (res * x) % mod;
        y = y >> 1;
        x = (x * x) % mod;
    }
    return res;
}

// Java code to implement the approach

import java.io.*;
import java.util.*;

// Driver Class
public class Main {
    // Mod value
    static int mod = 998244353;

    // Driver Function
    public static void main(String[] args)
        throws IOException
    {

        // Inputs
        int N = 4;
        int[] arr = { 3, 1, 4, 2 };

        // Function_call
        Find_Distinct_Arrays(N, arr);
    }

    // Function to output array B[]
    public static void Find_Distinct_Arrays(int N,
                                            int[] arr)
    {
        // Array to store the count of elements which have
        // both next greater and prev greater within the
        // current range
        int[] have = new int[N];
        int cnt = 0;

        // Arrays to store previous and next greater
        // elements respectively
        int[] prev = PrevGreater(arr);
        int[] next = NextGreater(arr);

        for (int i = 0; i < N; i++) {
            // If the current element has both
            // a previous and a next greater element
            if (prev[i] != -1 && next[i] != -1) {
                // Increment the count of elements which
                // have both
                // next greater and prev greater within the
                // current range
                // The count is stored at the index of the
                // next greater element
                have[next[i]]++;
            }
            // Add the count of such elements up to the
            // current index to the running total
            cnt += have[i];
            // Print 2 to the power of the running total,
            // modulo 998244353
            // This is the number of distinct arrays that
            // can be generated through the operations for
            // the current prefix of the permutation
            System.out.print(binpow(2, cnt) + " ");
        }
    }

    // Function to calculate prev greater element using
    // stack
    public static int[] PrevGreater(int[] arr)
    {
        // length of arr[]
        int N = arr.length;

        // Stack to find previous and next greater element
        Stack<Integer> s = new Stack<>();

        // Array to store the previous greater element
        int[] prev = new int[N];

        // Filling array with all -1s initially
        Arrays.fill(prev, -1);

        // Find previous greater element for each element
        for (int i = 0; i < N; i++) {
            while (!s.isEmpty() && arr[s.peek()] < arr[i]) {
                s.pop();
            }
            if (!s.isEmpty()) {
                prev[i] = s.peek();
            }
            s.push(i);
        }

        // returning array
        return prev;
    }

    // Function to calculate next greater element using
    // stack
    public static int[] NextGreater(int[] arr)
    {

        // length of arr[]
        int N = arr.length;

        // Stack to find previous and next greater element
        Stack<Integer> s = new Stack<>();

        // Array to store next greater element
        int[] next = new int[N];

        // Filling array with all -1s initially
        Arrays.fill(next, -1);

        // Find next greater element for each element
        for (int i = N - 1; i >= 0; i--) {
            while (!s.isEmpty() && arr[s.peek()] < arr[i]) {
                s.pop();
            }
            if (!s.isEmpty()) {
                next[i] = s.peek();
            }
            s.push(i);
        }

        // returning array
        return next;
    }

    // Function to calculate power in log(n) time
    static int binpow(int x, int y)
    {
        int res = 1;
        while (y > 0) {
            if (y % 2 == 1)
                res = (res * x) % mod;
            y = y >> 1;
            x = (x * x) % mod;
        }
        return res;
    }
}

from __future__ import print_function


def find_distinct_arrays(N, arr):
    # Array to store the count of elements which have
    # both next greater and prev greater within the
    # current range
    have = [0] * N
    cnt = 0

    # Arrays to store previous and next greater
    # elements respectively
    prev = prev_greater(arr)
    next = next_greater(arr)

    for i in range(N):
        # If the current element has both
        # a previous and a next greater element
        if prev[i] != -1 and next[i] != -1:
            # Increment the count of elements which
            # have both next greater and prev greater within the
            # current range
            # The count is stored at the index of the
            # next greater element
            have[next[i]] += 1
        # Add the count of such elements up to the
        # current index to the running total
        cnt += have[i]
        # Print 2 to the power of the running total,
        # modulo 998244353
        # This is the number of distinct arrays that
        # can be generated through the operations for
        # the current prefix of the permutation
        # End with a space instead of newline
        print(binpow(2, cnt) % 998244353),

# Function to calculate prev greater element using stack


def prev_greater(arr):
    # length of arr[]
    N = len(arr)

    # Stack to find previous greater element
    s = []

    # Array to store the previous greater element
    prev = [-1] * N

    # Find previous greater element for each element
    for i in range(N):
        while s and arr[s[-1]] < arr[i]:
            s.pop()
        if s:
            prev[i] = s[-1]
        s.append(i)

    # returning array
    return prev

# Function to calculate next greater element using stack


def next_greater(arr):
    # length of arr[]
    N = len(arr)

    # Stack to find next greater element
    s = []

    # Array to store next greater element
    next = [-1] * N

    # Find next greater element for each element
    for i in range(N - 1, -1, -1):
        while s and arr[s[-1]] < arr[i]:
            s.pop()
        if s:
            next[i] = s[-1]
        s.append(i)

    # returning array
    return next

# Function to calculate power in log(n) time


def binpow(x, y):
    res = 1
    while y > 0:
        if y % 2 == 1:
            res = (res * x) % 998244353
        y //= 2
        x = (x * x) % 998244353
    return res


# Driver code
if __name__ == "__main__":
    # Inputs
    N = 4
    arr = [3, 1, 4, 2]

    # Function call
    find_distinct_arrays(N, arr)

//C# code to implement the approach
using System;
using System.Collections.Generic;

class Program
{
    // Mod value
    static int mod = 998244353;

    // Driver Function
    static void Main()
    {
        // Inputs
        int N = 4;
        int[] arr = { 3, 1, 4, 2 };

        // Function call
        FindDistinctArrays(N, arr);
    }

    // Function to output array B[]
    static void FindDistinctArrays(int N, int[] arr)
    {
        // Array to store the count of elements which have
        // both next greater and prev greater within the
        // current range
        int[] have = new int[N];
        int cnt = 0;

        // Arrays to store previous and next greater
        // elements respectively
        int[] prev = PrevGreater(arr);
        int[] next = NextGreater(arr);

        for (int i = 0; i < N; i++)
        {
            // If the current element has both
            // a previous and a next greater element
            if (prev[i] != -1 && next[i] != -1)
            {
                // Increment the count of elements which
                // have both
                // next greater and prev greater within the
                // current range
                // The count is stored at the index of the
                // next greater element
                have[next[i]]++;
            }
            // Add the count of such elements up to the
            // current index to the running total
            cnt += have[i];
            // Print 2 to the power of the running total,
            // modulo 998244353
            // This is the number of distinct arrays that
            // can be generated through the operations for
            // the current prefix of the permutation
            Console.Write(BinPow(2, cnt) + " ");
        }
    }

    // Function to calculate prev greater element using stack
    static int[] PrevGreater(int[] arr)
    {
        // length of arr[]
        int N = arr.Length;

        // Stack to find previous and next greater element
        Stack<int> s = new Stack<int>();

        // Array to store the previous greater element
        int[] prev = new int[N];

        // Filling array with all -1s initially
        for (int i = 0; i < N; i++)
        {
            prev[i] = -1;
        }

        // Find previous greater element for each element
        for (int i = 0; i < N; i++)
        {
            while (s.Count > 0 && arr[s.Peek()] < arr[i])
            {
                s.Pop();
            }
            if (s.Count > 0)
            {
                prev[i] = s.Peek();
            }
            s.Push(i);
        }

        // returning array
        return prev;
    }

    // Function to calculate next greater element using stack
    static int[] NextGreater(int[] arr)
    {
        // length of arr[]
        int N = arr.Length;

        // Stack to find previous and next greater element
        Stack<int> s = new Stack<int>();

        // Array to store next greater element
        int[] next = new int[N];

        // Filling array with all -1s initially
        for (int i = 0; i < N; i++)
        {
            next[i] = -1;
        }

        // Find next greater element for each element
        for (int i = N - 1; i >= 0; i--)
        {
            while (s.Count > 0 && arr[s.Peek()] < arr[i])
            {
                s.Pop();
            }
            if (s.Count > 0)
            {
                next[i] = s.Peek();
            }
            s.Push(i);
        }

        // returning array
        return next;
    }

    // Function to calculate power in log(n) time
    static int BinPow(int x, int y)
    {
        int res = 1;
        while (y > 0)
        {
            if (y % 2 == 1)
                res = (res * x) % mod;
            y = y >> 1;
            x = (x * x) % mod;
        }
        return res;
    }
}

// Function to calculate previous greater element using stack
function prevGreater(arr) {
    const N = arr.length;
    const prev = new Array(N).fill(-1);
    const stack = [];

    for (let i = 0; i < N; i++) {
        while (stack.length && arr[stack[stack.length - 1]] < arr[i]) {
            stack.pop();
        }
        if (stack.length) {
            prev[i] = stack[stack.length - 1];
        }
        stack.push(i);
    }

    return prev;
}

// Function to calculate next greater element using stack
function nextGreater(arr) {
    const N = arr.length;
    const next = new Array(N).fill(-1);
    const stack = [];

    for (let i = N - 1; i >= 0; i--) {
        while (stack.length && arr[stack[stack.length - 1]] < arr[i]) {
            stack.pop();
        }
        if (stack.length) {
            next[i] = stack[stack.length - 1];
        }
        stack.push(i);
    }

    return next;
}

// Function to calculate power in log(n) time
function binpow(x, y) {
    let res = 1;
    while (y > 0) {
        if (y % 2 === 1) {
            res = (res * x) % 998244353;
        }
        y = Math.floor(y / 2);
        x = (x * x) % 998244353;
    }
    return res;
}

// Function to find distinct arrays
function findDistinctArrays(N, arr) {
    const have = new Array(N).fill(0);
    let cnt = 0;
    const prev = prevGreater(arr);
    const next = nextGreater(arr);

    for (let i = 0; i < N; i++) {
        if (prev[i] !== -1 && next[i] !== -1) {
            have[next[i]] += 1;
        }
        cnt += have[i];
        process.stdout.write(`${binpow(2, cnt) % 998244353} `);
    }
}

const N = 4;
const arr = [3, 1, 4, 2];
findDistinctArrays(N, arr);

1 1 2 2 

Time Complexity: O(N log N)
Auxiliary Space: O(N), As Stack and Prev[] and Next[] Arrays are used of Size N.