HackerRank: Friend Circles (DFS in Java)

Question

I got this problem on a HackerRank challenge and this was the best I could come up with. It's basically just using DFS to find connected components (in this case, friend circles). I was having a tough time figuring out how to track unvisited nodes. Please let me know how I can improve this code. It appears to work, but the test cases given were quite simple.

Problem:

There are N students in a class. Some of them are friends, while some are not. Their friendship is transitive in nature, i.e., if A is friend of B and B is friend of C, then A is also friend of C. A friend circle is a group of students who are directly or indirectly friends.

You have to complete a function int friendCircles(char[][] friends) which returns the number of friend circles in the class. Its argument, friends, is a NxN matrix which consists of characters "Y" or "N". If friends[i][j] == "Y" then i-th and j-th students are friends with each other, otherwise not. You have to return the total number of friend circles in the class.

Sample Input 0:
4
YYNN
YYYN
NYYN
NNNY
Sample Output 0:
2

Explanation 0:
There are two pairs of friends [0, 1] and [1, 2]. So [0, 2] is also a pair of friends by transitivity.
So first friend circle contains (0, 1, 2) and second friend circle contains only student 3

Sample Input 1:
5
YNNNN
NYNNN
NNYNN
NNNYN
NNNNY
Sample output 1:
5

Constraints (sorry, couldn't get formatting down so I had to put the constraints down here):

• 1 <= N <= 300.
• Each element of matrix friends will be "Y" or "N".
• Number of rows and columns will be equal in friends.
• friends[i][j] = "Y", where 0 <= i < N.
• friends[i][j] = friends[j][i], where 0 <= i < j < N.

Solution:

import java.util.ArrayDeque;
import java.util.Deque;
import java.util.HashSet;
import java.util.Set;

public class Solution {
    public static int friendsCircle(char[][] friends) {
        // The only alternative I could think of, instead of
        // tracking unvisited nodes, was to put visited nodes
        // in a set and then do setOfAllNodes.removeAll(visited)
        // to see which nodes are still unvisited
        Set<Integer> unvisited = new HashSet<>();
        boolean[] visited = new boolean[friends.length];
        Deque<Integer> stack = new ArrayDeque<>();
        int connectedComponents = 0;
        for (int i = 0; i < friends.length; i++) {
            unvisited.add(i);
        }
        // dfs on friends matrix
        while (!unvisited.isEmpty()) {
            stack.push(unvisited.iterator().next());
            connectedComponents++;
            while (!stack.isEmpty()) {
                int currVertex = stack.pop();
                if (visited[currVertex] == false) {
                    visited[currVertex] = true;
                    unvisited.remove(currVertex);
                    for (int i = 0; i < friends[currVertex].length; i++) {
                        if (friends[currVertex][i] == 'Y' && visited[i] == false) {
                            stack.push(i);
                        }
                    }
                }
            }
        }
        return connectedComponents;
    }

    public static void main(String[] args) {
        char[][] friends = {
                {'Y','Y','N','N'},
                {'Y','Y','Y','N'},
                {'N','Y','Y','N'},
                {'N','N','N','Y'}
        };
        System.out.println(friendsCircle(friends));
    }
}

Answer 1

Welcome to Code Review. Your code is good and easy to read, you can iterate over only half of the matrix (in my case the elements of the matrix a[i,j] with i < j) because from the test matrices always results a[i, j] = a[j, i] so friendship in symmetric . I'm using a TreeMap<Integer, Set<Integer>> to store initial situation to guarantee natural order of the keys so it will always the lower index the index where friends will be added : everybody is friend of himself (reflexive property) so map[i]={ i }.

int n = friends.length;
Map<Integer, Set<Integer>> map = new TreeMap<>();
for (int i = 0; i < n; ++i) {
    Set<Integer> set = new TreeSet<>();
    set.add(i);
    map.put(i, set);
}

Now I check all elements a[i, j] with i < j and see if i appears as key in the map: if yes I add j to map[j], otherwise it means that i appears in another set and I will add j to this set. At the end in any case I will remove the key j:

for (int i = 0; i < n; ++i) {
    for (int j = 0; j < n; ++j) {
        if (i < j && friends[i][j] == 'Y') {
            if (map.containsKey(i)) {
                map.get(i).add(j);      
            } else {
                for (Integer key : map.keySet()) {
                    Set<Integer> set = map.get(key);
                    if (set.contains(i)) {
                        set.add(j);
                    }
                }
            }
            map.remove(j);
        }
    }
}

The number of circles will coincide with the number of keys present in the map at the end:

public static int CountFriendsCircles(char[][] friends) {
    int n = friends.length;
    Map<Integer, Set<Integer>> map = new TreeMap<>();
    for (int i = 0; i < n; ++i) {
        Set<Integer> set = new TreeSet<>();
        set.add(i);
        map.put(i, set);
    }

    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < n; ++j) {
            if (i < j && friends[i][j] == 'Y') {
                if (map.containsKey(i)) {
                    map.get(i).add(j);      
                } else {
                    for (Integer key : map.keySet()) {
                        Set<Integer> set = map.get(key);
                        if (set.contains(i)) {
                            set.add(j);
                        }
                    }
                }
                map.remove(j);
            }
        }
    }

    return map.size();
}

Answer 2

I have suggestions for you.

1. The unvisited should be renamed unvisitedIndexes in my opinion.

2. I suggest that you create a method to build the unvisited set to separate the logic from the main method.

public static int friendsCircle(char[][] friends) {
   Set<Integer> unvisited = buildUnvisitedSet(friends);
}

private static Set<Integer> buildUnvisitedSet(char[][] friends) {
   Set<Integer> unvisited = new HashSet<>();
   for (int i = 0; i < friends.length; i++) {
      unvisited.add(i);
   }
   return unvisited;
}

3. Instead of using visited[currVertex] == false and visited[i] == false, you can use !visited[currVertex] and !visited[i].

4. You can make a method to check if the current node is visited or not, instead of handling the array each time.

public static int friendsCircle(char[][] friends) {
   //[...]
   if (isNotVisited(visited, currVertex)) {
      //[...]
   }
   //[...]
}

public static boolean isNotVisited(boolean[] visited, int index) {
   return !visited[index];
}

5. I suggest that you extract the logic that finds the friend position in a method, this will make the code shorter and easier to read.

public static int friendsCircle(char[][] friends) {
   //[...]
   int friendPosition = findFriendPosition(friends, visited, currVertex);
   if (friendPosition != -1) {
      stack.push(friendPosition);
   }
   //[...]
}

private static int findFriendPosition(char[][] friends, boolean[] visited, int currVertex) {
   for (int i = 0; i < friends[currVertex].length; i++) {
      if (friends[currVertex][i] == 'Y' && isNotVisited(visited, i)) {
         return i;
      }
   }
   return -1;
}

Refoctored code

public static void main(String[] args) {
   char[][] friends = {
      {'Y', 'Y', 'N', 'N'},
      {'Y', 'Y', 'Y', 'N'},
      {'N', 'Y', 'Y', 'N'},
      {'N', 'N', 'N', 'Y'}
   };
   System.out.println(friendsCircle(friends));
}

public static int friendsCircle(char[][] friends) {

   boolean[] visited = new boolean[friends.length];
   Set<Integer> unvisitedIndexes = buildUnvisitedSet(friends);

   Deque<Integer> stack = new ArrayDeque<>();
   int connectedComponents = 0;

   // dfs on friends matrix
   while (!unvisitedIndexes.isEmpty()) {
      stack.push(unvisitedIndexes.iterator().next());
      connectedComponents++;

      while (!stack.isEmpty()) {
         int currVertex = stack.pop();

         if (isNotVisited(visited, currVertex)) {
            visited[currVertex] = true;
            unvisitedIndexes.remove(currVertex);

            int friendPosition = findFriendPosition(friends, visited, currVertex);
            if (friendPosition != -1) {
               stack.push(friendPosition);
            }
         }
      }
   }

   return connectedComponents;
}

private static int findFriendPosition(char[][] friends, boolean[] visited, int currVertex) {
   for (int i = 0; i < friends[currVertex].length; i++) {
      if (friends[currVertex][i] == 'Y' && isNotVisited(visited, i)) {
         return i;
      }
   }
   return -1;
}

public static boolean isNotVisited(boolean[] visited, int index) {
   return !visited[index];
}

private static Set<Integer> buildUnvisitedSet(char[][] friends) {
   // The only alternative I could think of, instead of
   // tracking unvisited nodes, was to put visited nodes
   // in a set and then do setOfAllNodes.removeAll(visited)
   // to see which nodes are still unvisited
   Set<Integer> unvisited = new HashSet<>();
   for (int i = 0; i < friends.length; i++) {
      unvisited.add(i);
   }
   return unvisited;
}