Falling Dominoes using stack in O(n) similar to next greater element for each element


#1

Code:-

public class Solution {
static class Node{
    int id;
    int l;
    int r;
    int ans=-1;
    
    Node(int id1,int l1,int r1){
        id=id1;
        l=l1;
        r=r1;
    }
    
}
public int[] solve(int[] A, int[] B) {
    
    Node nodes[]=new Node[A.length];
    for(int i=0;i<A.length;i++){
        nodes[i]=new Node(i,A[i],A[i]+B[i]);
    }
    
    
    Arrays.sort(nodes,new Comparator<Node>(){
        
        public int compare(Node n1,Node n2){
            return Integer.compare(n1.l,n2.l);
        }
    });
    
    
    int n=A.length;
    
    //Using Stack in extra space
    
    Stack<Integer> s=new Stack();
    //logic similar to next greater element of an array
    
    for(int i=0;i<n;i++){
       Node node=nodes[i];
       while(true){
           if(s.isEmpty()){
               s.push(i);
               break;
           }else if(nodes[s.peek()].r<=node.l){
               int j=s.pop();
               nodes[j].ans=i-j;
           }else{
               s.push(i);
               break;
           }
           
       }    
    }
    
    
    while(!s.isEmpty()){
        int index=s.pop();
        nodes[index].ans=n-index;
    }
    
    
    Arrays.sort(nodes,new Comparator<Node>(){
        
        public int compare(Node n1,Node n2){
            return Integer.compare(n1.id,n2.id);
        }
    });
    
    
    int ans[]=new int[n];
    
    for(int i=0;i<ans.length;i++){
        ans[i]=nodes[i].ans;
    }
    
    
    return ans;
    
    
}

}