Count pair of indices in Array having equal Prefix-MEX and Suffix-MEX

Given an array arr[] of N elements, the task is to find the number of pairs of indices such that the Prefix-MEX of the one index is equal to the Suffix-MEX of the other index. Order of indices in pair matters.

MEX of an array refers to the smallest missing non-negative integer of the array. For the given problems (i, j) and (j, i) are considered different pairs.

Examples:

Input: A[] = {1, 0, 2, 0, 1}
Output: 11
Explanation: In array A,   Prefix-MEX for each element are as {0, 2, 3, 3, 3}, similarly, Suffix-MEX for each element are {3, 3, 3, 2, 0} and all possible pairs of indexes such that prefix-MEX of one is equal to Suffix-MEX of other are (1-based indexing):
(1, 5): prefix of 1 is equal to the suffix of 5
(2, 4):  prefix of 2 is equal to the suffix of 4
(3, 3):  prefix of 3 is equal to the suffix of 3
(3, 2):  prefix of 3 is equal to the suffix of 2
(3, 1):  prefix of 3 is equal to the suffix of 1
(4, 3):  prefix of 4 is equal to the suffix of 3
(4, 2):  prefix of 4 is equal to the suffix of 2
(4, 1):  prefix of 4 is equal to the suffix of 1
(5, 3):  prefix of 5 is equal to the suffix of 3
(5, 2):  prefix of 5 is equal to the suffix of 2
(5, 1):  prefix of 5 is equal to the suffix of 1

Input: A[] = {1, 2, 3}
Output: 9

Naive Approach: 

The most straightforward way will be to find each element’s Prefix–MEX and suffix–MEX and count a total number of possible pairs and increase the count. And for each element, we have to traverse the complete array again and again therefore complexity will be of order N^2.

Time Complexity: O(N^2)
Auxiliary Space: O(1)

Efficient Approach: In this approach, we will pre-compute and store Prefix-MEX and Suffix-MEX for each element.

Follow the steps below to calculate the Prefix-MEX array for the given array: 

• Find the maximum element in the given array arr.
• Create a set of integers and store the numbers from 0 to the max element in the set.
• Traverse through the array from i = 0 to N-1:
  • For each element, erase that element from the set.
  • Now find the smallest element remaining in the set.
  • Store the value of this smallest element remaining in the set in the resultant array.
• Return the resultant array as the required answer.

Follow the steps below to implement the idea:

• Create a variable count and initialize it to 0 to store the number of indexes.
• Calculate a Prefix-MEX for arr as P for a given array.
• Reverse the given array arr and calculate the Prefix-MEX array for this revered array as S and reverse S.
• Create a map mp to count the frequency of each element in P.
• Traverse through P increment frequency of current element.
• Traverse through S and check if the current element exists in mp.
  • if the current element does not exist in mp then continue.
  • else add the frequency of the current element in the variable count.
• print the value of count as a required result.

Below is the implementation of the above approach:

C++

// C++ code to implement the approach

#include <bits/stdc++.h>
using namespace std;

// Function to find the prefix MEX
// for each array element taking vector
// and its size as parameter
vector<int> Prefix_MEX(vector<int>& A, int n)
{

    // Maximum element in vector A
    int mx_element = *max_element(A.begin(), A.end());

    // Set to store all elements for
    // 0 to mx_element
    set<int> s;

    // Vector to store Prefix-MEX for
    // given input array
    vector<int> B(n);

    // Store all number from 0
    // to maximum element + 1 in a set
    for (int i = 0; i <= mx_element + 1; i++) {
        // inserting elements in set
        s.insert(i);
    }

    // Loop to calculate MEX for each
    // index in the array
    for (int i = 0; i < n; i++) {

        // Checking if A[i] is present
        // in set
        auto it = s.find(A[i]);

        // If present then we erase
        // that element
        if (it != s.end())
            s.erase(it);

        // Store the first element of set
        // in vector B as Mex of
        // prefix vector
      //cout<<*s.begin()<<" ";
        B[i] = *s.begin();
    }

    // Return the vector B
    return B;
}

void countPairs(vector<int>& arr, int N)
{

    // Count variable to store number of
    // indices whose prefix-mex and
    // suffix-mex are equal
    int count = 0;

    // Vector P stores the Prefix-MEX
    // for given array
    vector<int> P = Prefix_MEX(arr, N);

    // Reversing given array
    reverse(arr.begin(), arr.end());

    // vector S stores the reverse of
    // Suffix-MEX for given array

    vector<int> S = Prefix_MEX(arr, N);

    // Reversing vector S will give
    // suffix-MEX for given array
    reverse(S.begin(), S.end());

    // map to count frequency of each
    // element in arr1
    map<int, int> mp;

    // Counting frequency of each
    // element of arr1
    for (int i = 0; i < N; i++) {
        mp[P[i]]++;
    }

    // Traversal through arr S
    for (int i = 0; i < N; i++) {

        // Check if element does not exist
        // in map
        if (mp.find(S[i]) == mp.end()) {
            continue;
        }

        // if exist add it's frequency
        // to count
        else {
            count += mp[S[i]];
        }
    }
    cout << count << endl;
}

// Driver code
int main()
{
    vector<int> arr = { 1, 0, 2, 0, 1 };
    int N = arr.size();

    // Function Call
    countPairs(arr, N);
    arr = { 1, 2, 3 };
    N = arr.size();

    // Function Call
    countPairs(arr, N);

    return 0;
}

Java

import java.util.*;

public class Main {
  static List<Integer> prefixMex(List<Integer> A) {
    int mxElement = Collections.max(A);

    Set<Integer> s = new HashSet<>();
    for (int i = 0; i < mxElement + 2; i++) {
      s.add(i);
    }

    List<Integer> B = new ArrayList<>(Collections.nCopies(A.size(), 0));

    for (int i = 0; i < A.size(); i++) {
      s.remove(A.get(i));
      B.set(i, Collections.min(s));
    }

    return B;
  }

  static int countPairs(List<Integer> arr) {
    int n = arr.size();

    List<Integer> P = prefixMex(arr);
    List<Integer> S = new ArrayList<>(arr);
    Collections.reverse(S);
    S = prefixMex(S);
    Collections.reverse(S);

    Map<Integer, Integer> mp = new HashMap<>();
    for (int p : P) {
      mp.put(p, mp.getOrDefault(p, 0) + 1);
    }

    int count = 0;
    for (int s : S) {
      count += mp.getOrDefault(s, 0);
    }

    return count;
  }

  public static void main(String[] args) {
    List<Integer> arr = Arrays.asList(1, 0, 2, 0, 1);
    System.out.println(countPairs(arr));

    arr = Arrays.asList(1, 2, 3);
    System.out.println(countPairs(arr));
  }
}

Python3

from typing import List, Tuple

def prefix_mex(A: List[int]) -> List[int]:
    mx_element = max(A)

    s = set(range(mx_element + 2))

    B = [0] * len(A)

    for i, a in enumerate(A):
        s.discard(a)
        B[i] = min(s)

    return B

def count_pairs(arr: List[int]) -> int:
    n = len(arr)

    P = prefix_mex(arr)
    S = prefix_mex(arr[::-1])[::-1]

    mp = {}
    for p in P:
        mp[p] = mp.get(p, 0) + 1

    count = 0
    for s in S:
        count += mp.get(s, 0)

    return count

arr = [1, 0, 2, 0, 1]
print(count_pairs(arr))

arr = [1, 2, 3]
print(count_pairs(arr))

C#

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

public class GFG {

    // Function to find the prefix MEX
    // for each array element taking vector
    // and its size as parameter
    static List<int> Prefix_MEX(List<int> A, int n)
    {
        // Maximum element in List A
        int mxElement = A.Max();

        // Set to store all elements for
        // 0 to mx_element
        HashSet<int> s = new HashSet<int>();

        // Store all number from 0
        // to maximum element + 1 in a set
        for (int i = 0; i < mxElement + 2; i++) {
            // inserting elements in set
            s.Add(i);
        }

        // List to store Prefix-MEX for
        // given input array
        List<int> B = new List<int>(new int[A.Count]);

        // Loop to calculate MEX for each
        // index in the array
        for (int i = 0; i < n; i++) {
            // Checking if A[i] is present in set
            // If present then we remove that element
            s.Remove(A[i]);

            // Store the first element of set
            // in vector B as Mex of
            B[i] = s.Min();
        }

        return B;
    }

    static int CountPairs(List<int> arr, int n)
    {

        // Vector P stores the Prefix-MEX
        // for given array
        List<int> P = Prefix_MEX(arr, n);

        // vector S stores the reverse of
        // Suffix-MEX for given array
        List<int> S = new List<int>(arr);

        // Reversing given array
        S.Reverse();
        S = Prefix_MEX(S, n);

        // Reversing vector S will give
        // suffix-MEX for given array
        S.Reverse();

         // map to count frequency of each
        // element in arr1
        Dictionary<int, int> mp
            = new Dictionary<int, int>();

        // Counting frequency of each
        // element of arr1
        foreach(int p in P) {
            if (mp.ContainsKey(p)) {
                mp[p]++;
            }
            else {
                mp[p] = 1;
            }
        }

        // Count variable to store number of
        // indices whose prefix-mex and
        // suffix-mex are equal
        int count = 0;

        // Traversal through arr S
        foreach(int s in S) {
            if (mp.ContainsKey(s)) {
                count += mp[s];
            }
        }
        return count;
    }

    // Driver code
    public static void Main(string[] args)
    {
        List<int> arr = new List<int>{ 1, 0, 2, 0, 1 };
        int N = arr.Count;

        // Function Call
        Console.WriteLine(CountPairs(arr, N));

        arr = new List<int>{ 1, 2, 3 };
        N = arr.Count;

        // Function Call
        Console.WriteLine(CountPairs(arr, N));
    }
}

// This Code is Contributed by Prasad Kandekar(prasad264)

Javascript

// javascript code to implement the approach

// Function to find the prefix MEX
// for each array element taking vector
// and its size as parameter
function Prefix_MEX(A, n)
{

    // Maximum element in vector A
    let mx_element = Number.MIN_SAFE_INTEGER;
    for(let i=0; i<A.length; i++)
        mx_element=Math.max(mx_element, A[i]);

    // Set to store all elements for
    // 0 to mx_element
    let s=new Set();

    // Vector to store Prefix-MEX for
    // given input array
    let B=new Array(n);

    // Store all number from 0
    // to maximum element + 1 in a set
    for (let i = 0; i <= mx_element + 1; i++) {
        // inserting elements in set
        s.add(i);
    }

    // Loop to calculate MEX for each
    // index in the array
    for (let i = 0; i < n; i++) {

        // Checking if A[i] is present
        // in set
        //auto it = s.find(A[i]);

        // If present then we erase
        // that element
        if (s.has(A[i]))
            s.delete(A[i]);

        // Store the first element of set
        // in vector B as Mex of
        // prefix vector
        const [first]=s;
        B[i] = first;
    }

    // Return the vector B
    return B;
}

function countPairs(arr, N)
{

    // Count variable to store number of
    // indices whose prefix-mex and
    // suffix-mex are equal
    let count = 0;

    // Vector P stores the Prefix-MEX
    // for given array
    let P = Prefix_MEX(arr, N);

    // Reversing given array
    arr.reverse();
    // vector S stores the reverse of
    // Suffix-MEX for given array

    let S = Prefix_MEX(arr, N);

    // Reversing vector S will give
    // suffix-MEX for given array
    S.reverse();

    // map to count frequency of each
    // element in arr1
    let mp=new Map();

    // Counting frequency of each
    // element of arr1
    for (let i = 0; i < N; i++) {
        //mp[P[i]]++;
        if(mp.has(P[i]))
            mp.set(P[i], mp.get(P[i])+1);
        else
            mp.set(P[i],1);
    }

    // Traversal through arr S
    for (let i = 0; i < N; i++) {

        // Check if element does not exist
        // in map
        if (!mp.has(S[i])) {
            continue;
        }

        // if exist add it's frequency
        // to count
        else {
            count += mp.get(S[i]);
        }
    }
    console.log(count);
}

// Driver code
    let arr = [ 1, 0, 2, 0, 1 ];
    let N = arr.length;

    // Function Call
    countPairs(arr, N);
    arr = [ 1, 2, 3 ];
    N = arr.length;

    // Function Call
    countPairs(arr, N);

   // This code is contributed by poojaagarwal2.

Output

11
9

Time Complexity: O(N * log N), log N for inserting and deleting elements from the set, and O(N) for traversing the array.
Auxiliary Space:  O(N) for storing Prefix-MEX array.