import java.util.*;

public class GetSuggestedFriends<T> {

    private HashMap<T, ArrayList<T>> adj = new HashMap<>(); //graph
    private List<Set<T>> groups = new ArrayList<>();
    
    //Time O(1), Space O(1)
    public void addFriendship(T src, T dest) {		
		adj.putIfAbsent(src, new ArrayList<T>());
		adj.get(src).add(dest); 		
		adj.putIfAbsent(dest, new ArrayList<T>());
		adj.get(dest).add(src); 
    }
 
    //Find group by connections, DFS wrapper, Time O(V+E), Space O(V)
    //V is total number of people, E is number of connections
    private void findGroups() {
    	Map<T, Boolean> visited = new HashMap<>();
    	for (T t: adj.keySet())
    	   visited.put(t, false);
        for (T t:adj.keySet()) {
            if (!visited.get(t)) {
            	Set<T> group = new HashSet<>();
                dfs(t, visited, group);
                groups.add(group);
            }
        }
    }
    
    //DFS + memoization, Time O(V+E), Space O(V) 
    private void dfs(T v, Map<T, Boolean> visited, Set<T> group ) {
        visited.put(v,true);
        group.add(v);
        for (T x : adj.get(v)) {
            if (!visited.get(x)) 
                dfs(x, visited, group);
        }
    }    
    
    //Find suggest friends, Time O(V+E), Space O(V)
	public Set<T> getSuggestedFriends (T a) {
		if (groups.isEmpty())
			findGroups();
		Set<T> res = new HashSet<>();
		for (Set<T> t : groups) {
			if (t.contains(a)) {
				res = t;
				break;
			}
		}
		if (res.size() > 0) //remove himself from result
			res.remove(a);
        for (T t : adj.get(a)) //remove already known friends from result
            res.remove(t);
		return res;
	}

	public static void main(String[] args) {	
		
		GetSuggestedFriends<String> g = new GetSuggestedFriends<>();
		g.addFriendship("Amy", "Chris");
		g.addFriendship("Sarah", "Joshua");
		g.addFriendship("Joshua", "Zoe");
		g.addFriendship("Sarah", "Jess");
		g.addFriendship("Amy", "Sam");
	
		String name = "Zoe";
		System.out.println("Suggestion friends of " + name + ": " + g.getSuggestedFriends(name));
	
		name = "Sam";
		System.out.println("Suggestion friends of " + name + ": " + g.getSuggestedFriends(name));
	}
}