Maximum sum of segments among all segments formed in array after Q queries

Approach

We can solve the given problem using Disjoint-Set Union. Here in the given problem, after every query, we break the array into many small sets and then calculate the maximum sum over all segments. When we break a set into two smaller sets, the total sum is also divided. We can see that a query divides a set into at max two further sets, and their total sum is also divided accordingly. For example, let’s say there is a set S = {a1, a2, q, b1, b2} and currently we need to process a ‘query[i] = q’. Then after this query set S will be divided into two S1 = {a1, a2} and S2 = {b1, b2}.

The idea is to simulate this procedure in reverse order(i.e., instead of breaking the array into further sets after every query, we start from the last query and merge elements of the array). Initially, we put every element into different sets. Now we process queries in reverse order. For each query, do the current element’s union with its left & right element in the array (i.e., left & right sets). And simultaneously keep track of the sum of each set(using disjoin-set union). 

Algorithm

1. Declare vector setSum(N+1, 0), finalAns, and initialize a DSU class object of size N.

2. Initialize setSum[i] = a[i-1] (i.e, we are putting each element in a different set).

3. Insert ‘0’ in finalAns because, after the last query, all elements will be removed, and hence answer for the last query will be zero.

4. Now, iterate the queries in reverse order(using variable ‘i’) and do the following.
a) If p[query[i]] == -1, then set it as query[i] (i.e, it the current element is not part of any set, make it an independent set).
b) If (query[i]-1 >= 0) and (p[query[i]-1] != -1) then call Union(query[i], query[i]-1). (i.e, if element left to current query element exist and is part of some set S, then combine current query element with set S)
c) If (query[i]+1 >= 0) and (p[query[i]+1] != -1) then call Union(query[i], query[i]+1). (i.e, if element right to current query element exist and is part of some set S, then combine current query element with set S)
d) Update maxSegSum as max(maxSegSum, setSum[Find(query[i])]) and push it into finalAns.
5. Finally, reverse finalAns and print it.

Program

#include <bits/stdc++.h>
using namespace std;
 
#define vi vector<int>
 
 
class DSU{
public:
   
    int N; // 1 indexed
    vi p; // p[i] // Stores parent of i node.
    vi sz; // sz[i] // Stores the size of subtree of node i
 
    DSU(int inputN){
        N = inputN;
        // Initialize all elements as inidividual sets.
        p.resize(N+1, -1);
 
        // Initially size of each set is 1.
        sz.resize(N+1, 1);
    }
 
    // This function perform 'FIND' operation of Disjoint-Set Union data structure.
    int Find(int i){
        if(p[i] == i) return i;
        return p[i] = Find(p[i]);
    }
 
    // This function perform 'UNION' operation of Disjoint-Set Union data structure.
    void Union(int pu, int pv, vi& setSum){
        // Finding the root of respective sets.
        pu = Find(pu);
        pv = Find(pv);
       
        // Return if both elements belongs
        // to the same set.
        if(pu == pv) return;
 
        if(sz[pu] < sz[pv]) swap(pu, pv);
        // Now pu is the bigger component.
 
        // Update the parent of smaller set.
        p[pv] = pu;
 
        // Increment the size of bigger set as we 
        // are adding a smaller set in it.
        sz[pu] += sz[pv];
 
        // Update the sum of new set.
        setSum[pu] += setSum[pv];
    }
};
 
void solve(int N, vi& a, vi& query){
    // DSU object (has Union & Find methods).
    DSU dsu(N);
 
    // This vector stores of the sum of segments.
    vi setSum(N+1, 0);
 
    // This vector stores the maximum sum for each query.
    vi finalAns;
 
    // Initially every individual element is a set.
    for(int i=1; i<=N; i++) {
        setSum[i] = a[i-1];
    }
 
    // After processing all queries only empty sets will remain.
    // i.e. maximum sum will be zero.
    finalAns.push_back(0);
 
    // For storing maximum segment sum till current query.
    int maxSegSum = INT_MIN;
 
    for(int i=N-1; i>0; i--){
 
        /* If the current element isn't in any sets,
        set its parent as current element of the query.
        (i.e, make it a independent set.)*/
        if(dsu.p[query[i]] == -1){
            dsu.p[query[i]] = query[i];
        }
 
        /* If element left of query[i] in array a[]
        has been added to some set or is a set itself,
        Union it with current set */
        if(query[i]-1 >= 0 && dsu.p[query[i]-1] != -1){
            dsu.Union(query[i], query[i]-1, setSum);
        }
   
        /* If element right of query[i] in array a[]
        has been added to some set or is a set itself,
        Union it with current set */
        if(query[i]+1 <= N && dsu.p[query[i]+1] != -1){
            dsu.Union(query[i], query[i]+1, setSum);
        }

        // Updating the maxSegSum.
        maxSegSum = max(maxSegSum, setSum[dsu.Find(query[i])]);
 
        // Push the answer for query[i-1].
        finalAns.push_back(maxSegSum);
    }
 
    // Reverse the finalAns because we processed queries in a reverse order.
    reverse(finalAns.begin(), finalAns.end());
 
    // Print the final answers.
    for(int x: finalAns)cout << x << " ";
 
    cout << endl;
}
 
// Driver Function.
int main(){
    int tt; cin >> tt;
    while(tt--){
        int N; cin >> N;
 
        vi a(N), query(N);
 
        for(int i=0; i<N; i++) cin >> a[i];
        for(int i=0; i<N; i++) cin >> query[i];
    
        solve(N, a, query);
    }
    return 0;
}

You can also try this code with Online C++ Compiler

Run Code

INPUT

2
4
1 3 2 5
3 4 1 2
5
1 2 3 4 5
4 2 3 5 1

OUTPUT

5 4 3 0
6 5 5 1 0

Time Complexity

The overall time complexity of this approach is O(N * log(N)).

Space Complexity

The auxiliary space complexity of the program is O(N).

FAQs

1. What are the applications of the Disjoint-Set Union data structure(also known as Union-Find)?
   It has many applications. It is used for cycle detection, keeping track of connected components in an undirected graph, and as a sub-routine in Kruskal’s algorithm to find a minimum spanning tree.
    
2. What is the worst-case time complexity of union operations when size and path compression optimizations are used in DSU?
   The worst-case time complexity of union operations, when done using size & path compression, is O(M * logN), where N is the number of nodes and M is the number of operations.
    
3. In which operation is path compression optimization applied (in the implementation of DSU)?
   Path compression is used in the Find operation. It does not affect union operations in any way.

Key Takeaways

Cheers if you reached here!! 

This article discussed an intriguing problem using the Disjoint-Set Union data structure, also known as Union-Find. DSU-based problems are sometimes simple to implement, but finding an efficient approach remains difficult.

Check out this problem - Reverse Nodes In K Group

Yet learning never stops, and there is a lot more to learn. So head over to our practice platform Coding Ninjas Studio to practice top problems, attempt mock tests, read interview experiences, and much more. Till then, Happy Learning!!