Problem

A teacher must divide a class of students into two teams to play dodgeball. Unfortunately, not all the kids get along, and several refuse to be put on the same team as that of their enemies.

Given an adjacency list of students and their enemies, write an algorithm that finds a satisfactory pair of teams, or returns False if none exists.

Examples

Example 1

1
2
3

Input: {0: [3], 1: [2], 2: [1, 4], 3: [0, 4, 5], 4: [2, 3], 5: [3]}
Output: [{0, 1, 4, 5}, {2, 3}] or False
Explanation: We can divide students into two teams where no enemies are on the same team. Team 1: {0, 1, 4, 5}, Team 2: {2, 3}. All enemy relationships are satisfied (enemies are on different teams).

Example 2

1
2
3

Input: {0: [3], 1: [2], 2: [1, 3, 4], 3: [0, 2, 4, 5], 4: [2, 3], 5: [3]}
Output: False
Explanation: It's impossible to divide students into two teams. For example, students 2, 3, and 4 form a triangle of enemies, making it impossible to place them in just two teams without conflicts.

Example 3

1
2
3

Input: {0: [1], 1: [0, 2], 2: [1]}
Output: [{0, 2}, {1}]
Explanation: Students 0 and 2 can be on the same team since they're not enemies, while student 1 (who is enemies with both) goes to the other team.

Example 4

1
2
3

Input: {0: [], 1: [], 2: []}
Output: [{0, 1, 2}, {}] or [{0}, {1, 2}] or any valid division
Explanation: No enemy relationships exist, so any division is valid. All students can be on the same team.

Example 5

1
2
3

Input: {}
Output: [set(), set()]
Explanation: Empty input results in two empty teams.

Solution

Method 1 - DFS Graph Coloring

Intuition

This problem is equivalent to determining if a graph is bipartite. We need to check if we can color the graph with exactly two colors such that no adjacent vertices (enemies) have the same color. Each color represents a team.

Approach

  1. Model the problem as an undirected graph where students are vertices and enemy relationships are edges
  2. Use DFS to attempt 2-coloring of the graph
  3. Start coloring from any unvisited node and assign it color 0 (team 1)
  4. For each neighbor, assign the opposite color and recursively check
  5. If we encounter a conflict (trying to assign a different color to an already colored node), return False
  6. Handle disconnected components by checking all unvisited nodes
  7. If successful, partition students based on their colors and return the teams

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

class Solution {
public:
    // Option 1: Using optional for clear success/failure indication
    optional<vector<vector<int>>> divideTeams(unordered_map<int, vector<int>>& enemies) {
        if (enemies.empty()) return vector<vector<int>>{{}, {}};
        
        unordered_map<int, int> colors;
        unordered_set<int> students;
        
        // Get all students
        for (auto& pair : enemies) {
            students.insert(pair.first);
            for (int enemy : pair.second) {
                students.insert(enemy);
            }
        }
        
        // Try to color the graph
        for (int student : students) {
            if (colors.find(student) == colors.end()) {
                if (!dfs(student, 0, enemies, colors)) {
                    return nullopt; // Clear indication of failure
                }
            }
        }
        
        // Build teams based on colors
        vector<int> team1, team2;
        for (int student : students) {
            if (colors[student] == 0) {
                team1.push_back(student);
            } else {
                team2.push_back(student);
            }
        }
        
        return vector<vector<int>>{team1, team2};
    }
    
    // Option 2: Using pair with boolean flag
    pair<bool, vector<vector<int>>> divideTeamsWithFlag(unordered_map<int, vector<int>>& enemies) {
        if (enemies.empty()) return {true, {{}, {}}};
        
        unordered_map<int, int> colors;
        unordered_set<int> students;
        
        // Get all students
        for (auto& pair : enemies) {
            students.insert(pair.first);
            for (int enemy : pair.second) {
                students.insert(enemy);
            }
        }
        
        // Try to color the graph
        for (int student : students) {
            if (colors.find(student) == colors.end()) {
                if (!dfs(student, 0, enemies, colors)) {
                    return {false, {}}; // Clear failure indication
                }
            }
        }
        
        // Build teams based on colors
        vector<int> team1, team2;
        for (int student : students) {
            if (colors[student] == 0) {
                team1.push_back(student);
            } else {
                team2.push_back(student);
            }
        }
        
        return {true, {team1, team2}};
    }
    
private:
    bool dfs(int node, int color, unordered_map<int, vector<int>>& enemies, 
             unordered_map<int, int>& colors) {
        colors[node] = color;
        
        if (enemies.find(node) != enemies.end()) {
            for (int enemy : enemies[node]) {
                if (colors.find(enemy) == colors.end()) {
                    if (!dfs(enemy, 1 - color, enemies, colors)) {
                        return false;
                    }
                } else if (colors[enemy] == color) {
                    return false; // Conflict found
                }
            }
        }
        
        return true;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

func divideTeams(enemies map[int][]int) [][]int {
    if len(enemies) == 0 {
        return [][]int{{}, {}}
    }
    
    colors := make(map[int]int)
    students := make(map[int]bool)
    
    // Get all students
    for student, enemyList := range enemies {
        students[student] = true
        for _, enemy := range enemyList {
            students[enemy] = true
        }
    }
    
    // Try to color the graph
    for student := range students {
        if _, visited := colors[student]; !visited {
            if !dfs(student, 0, enemies, colors) {
                return [][]int{} // Return empty to indicate False
            }
        }
    }
    
    // Build teams based on colors
    var team1, team2 []int
    for student := range students {
        if colors[student] == 0 {
            team1 = append(team1, student)
        } else {
            team2 = append(team2, student)
        }
    }
    
    return [][]int{team1, team2}
}

func dfs(node, color int, enemies map[int][]int, colors map[int]int) bool {
    colors[node] = color
    
    if enemyList, exists := enemies[node]; exists {
        for _, enemy := range enemyList {
            if existingColor, visited := colors[enemy]; !visited {
                if !dfs(enemy, 1-color, enemies, colors) {
                    return false
                }
            } else if existingColor == color {
                return false // Conflict found
            }
        }
    }
    
    return true
}
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110

class Solution {
    // Option 1: Using Optional for clear success/failure indication
    public Optional<List<List<Integer>>> divideTeams(Map<Integer, List<Integer>> enemies) {
        if (enemies.isEmpty()) return Optional.of(Arrays.asList(new ArrayList<>(), new ArrayList<>()));
        
        Map<Integer, Integer> colors = new HashMap<>();
        Set<Integer> students = new HashSet<>();
        
        // Get all students
        for (Map.Entry<Integer, List<Integer>> entry : enemies.entrySet()) {
            students.add(entry.getKey());
            for (int enemy : entry.getValue()) {
                students.add(enemy);
            }
        }
        
        // Try to color the graph
        for (int student : students) {
            if (!colors.containsKey(student)) {
                if (!dfs(student, 0, enemies, colors)) {
                    return Optional.empty(); // Clear indication of failure
                }
            }
        }
        
        // Build teams based on colors
        List<Integer> team1 = new ArrayList<>();
        List<Integer> team2 = new ArrayList<>();
        
        for (int student : students) {
            if (colors.get(student) == 0) {
                team1.add(student);
            } else {
                team2.add(student);
            }
        }
        
        return Optional.of(Arrays.asList(team1, team2));
    }
    
    // Option 2: Using custom result class
    public static class DivisionResult {
        public final boolean possible;
        public final List<List<Integer>> teams;
        
        public DivisionResult(boolean possible, List<List<Integer>> teams) {
            this.possible = possible;
            this.teams = teams;
        }
    }
    
    public DivisionResult divideTeamsWithResult(Map<Integer, List<Integer>> enemies) {
        if (enemies.isEmpty()) {
            return new DivisionResult(true, Arrays.asList(new ArrayList<>(), new ArrayList<>()));
        }
        
        Map<Integer, Integer> colors = new HashMap<>();
        Set<Integer> students = new HashSet<>();
        
        // Get all students
        for (Map.Entry<Integer, List<Integer>> entry : enemies.entrySet()) {
            students.add(entry.getKey());
            for (int enemy : entry.getValue()) {
                students.add(enemy);
            }
        }
        
        // Try to color the graph
        for (int student : students) {
            if (!colors.containsKey(student)) {
                if (!dfs(student, 0, enemies, colors)) {
                    return new DivisionResult(false, new ArrayList<>()); // Clear failure
                }
            }
        }
        
        // Build teams based on colors
        List<Integer> team1 = new ArrayList<>();
        List<Integer> team2 = new ArrayList<>();
        
        for (int student : students) {
            if (colors.get(student) == 0) {
                team1.add(student);
            } else {
                team2.add(student);
            }
        }
        
        return new DivisionResult(true, Arrays.asList(team1, team2));
    }
    
    private boolean dfs(int node, int color, Map<Integer, List<Integer>> enemies, 
                       Map<Integer, Integer> colors) {
        colors.put(node, color);
        
        if (enemies.containsKey(node)) {
            for (int enemy : enemies.get(node)) {
                if (!colors.containsKey(enemy)) {
                    if (!dfs(enemy, 1 - color, enemies, colors)) {
                        return false;
                    }
                } else if (colors.get(enemy) == color) {
                    return false; // Conflict found
                }
            }
        }
        
        return true;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

class Solution:
    def divide_teams(self, enemies: Dict[int, List[int]]) -> Union[List[List[int]], bool]:
        if not enemies:
            return [[], []]
        
        colors = {}
        students = set()
        
        # Get all students
        for student, enemy_list in enemies.items():
            students.add(student)
            for enemy in enemy_list:
                students.add(enemy)
        
        # Try to color the graph
        for student in students:
            if student not in colors:
                if not self._dfs(student, 0, enemies, colors):
                    return False
        
        # Build teams based on colors
        team1 = [student for student in students if colors[student] == 0]
        team2 = [student for student in students if colors[student] == 1]
        
        return [team1, team2]
    
    def _dfs(self, node: int, color: int, enemies: Dict[int, List[int]], 
             colors: Dict[int, int]) -> bool:
        colors[node] = color
        
        if node in enemies:
            for enemy in enemies[node]:
                if enemy not in colors:
                    if not self._dfs(enemy, 1 - color, enemies, colors):
                        return False
                elif colors[enemy] == color:
                    return False  # Conflict found
        
        return True

Complexity

  • ⏰ Time complexity: O(V + E), where V is the number of students and E is the number of enemy relationships. We visit each student once and traverse each edge once during DFS
  • 🧺 Space complexity: O(V + E), for storing the color mapping, student set, and recursion stack which can go up to V levels deep

Method 2 - BFS Graph Coloring

Intuition

Instead of using DFS, we can use BFS to perform the 2-coloring. This approach uses a queue to level-order traverse the graph while maintaining the color constraints, which can be more intuitive for some problems.

Approach

  1. Create a queue and color mapping to track student assignments
  2. For each unvisited student, start BFS with initial color 0
  3. Process each student in the queue and assign opposite colors to their enemies
  4. If we find a conflict (enemy already has the same color), return False
  5. Continue until all connected components are processed
  6. Build the final teams based on color assignments

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

class Solution {
public:
    vector<vector<int>> divideTeams(unordered_map<int, vector<int>>& enemies) {
        if (enemies.empty()) return {{}, {}};
        
        unordered_map<int, int> colors;
        unordered_set<int> students;
        
        // Get all students
        for (auto& pair : enemies) {
            students.insert(pair.first);
            for (int enemy : pair.second) {
                students.insert(enemy);
            }
        }
        
        // Try to color using BFS
        for (int student : students) {
            if (colors.find(student) == colors.end()) {
                if (!bfs(student, enemies, colors)) {
                    return {}; // Return empty to indicate False
                }
            }
        }
        
        // Build teams
        vector<int> team1, team2;
        for (int student : students) {
            if (colors[student] == 0) {
                team1.push_back(student);
            } else {
                team2.push_back(student);
            }
        }
        
        return {team1, team2};
    }
    
private:
    bool bfs(int start, unordered_map<int, vector<int>>& enemies, 
             unordered_map<int, int>& colors) {
        queue<int> q;
        q.push(start);
        colors[start] = 0;
        
        while (!q.empty()) {
            int curr = q.front();
            q.pop();
            
            if (enemies.find(curr) != enemies.end()) {
                for (int enemy : enemies[curr]) {
                    if (colors.find(enemy) == colors.end()) {
                        colors[enemy] = 1 - colors[curr];
                        q.push(enemy);
                    } else if (colors[enemy] == colors[curr]) {
                        return false; // Conflict found
                    }
                }
            }
        }
        
        return true;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

func divideTeamsBFS(enemies map[int][]int) [][]int {
    if len(enemies) == 0 {
        return [][]int{{}, {}}
    }
    
    colors := make(map[int]int)
    students := make(map[int]bool)
    
    // Get all students
    for student, enemyList := range enemies {
        students[student] = true
        for _, enemy := range enemyList {
            students[enemy] = true
        }
    }
    
    // Try to color using BFS
    for student := range students {
        if _, visited := colors[student]; !visited {
            if !bfs(student, enemies, colors) {
                return [][]int{} // Return empty to indicate False
            }
        }
    }
    
    // Build teams
    var team1, team2 []int
    for student := range students {
        if colors[student] == 0 {
            team1 = append(team1, student)
        } else {
            team2 = append(team2, student)
        }
    }
    
    return [][]int{team1, team2}
}

func bfs(start int, enemies map[int][]int, colors map[int]int) bool {
    queue := []int{start}
    colors[start] = 0
    
    for len(queue) > 0 {
        curr := queue[0]
        queue = queue[1:]
        
        if enemyList, exists := enemies[curr]; exists {
            for _, enemy := range enemyList {
                if _, visited := colors[enemy]; !visited {
                    colors[enemy] = 1 - colors[curr]
                    queue = append(queue, enemy)
                } else if colors[enemy] == colors[curr] {
                    return false // Conflict found
                }
            }
        }
    }
    
    return true
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

class Solution {
    public List<List<Integer>> divideTeams(Map<Integer, List<Integer>> enemies) {
        if (enemies.isEmpty()) return Arrays.asList(new ArrayList<>(), new ArrayList<>());
        
        Map<Integer, Integer> colors = new HashMap<>();
        Set<Integer> students = new HashSet<>();
        
        // Get all students
        for (Map.Entry<Integer, List<Integer>> entry : enemies.entrySet()) {
            students.add(entry.getKey());
            for (int enemy : entry.getValue()) {
                students.add(enemy);
            }
        }
        
        // Try to color using BFS
        for (int student : students) {
            if (!colors.containsKey(student)) {
                if (!bfs(student, enemies, colors)) {
                    return new ArrayList<>(); // Return empty to indicate False
                }
            }
        }
        
        // Build teams
        List<Integer> team1 = new ArrayList<>();
        List<Integer> team2 = new ArrayList<>();
        
        for (int student : students) {
            if (colors.get(student) == 0) {
                team1.add(student);
            } else {
                team2.add(student);
            }
        }
        
        return Arrays.asList(team1, team2);
    }
    
    private boolean bfs(int start, Map<Integer, List<Integer>> enemies, 
                       Map<Integer, Integer> colors) {
        Queue<Integer> queue = new LinkedList<>();
        queue.offer(start);
        colors.put(start, 0);
        
        while (!queue.isEmpty()) {
            int curr = queue.poll();
            
            if (enemies.containsKey(curr)) {
                for (int enemy : enemies.get(curr)) {
                    if (!colors.containsKey(enemy)) {
                        colors.put(enemy, 1 - colors.get(curr));
                        queue.offer(enemy);
                    } else if (colors.get(enemy).equals(colors.get(curr))) {
                        return false; // Conflict found
                    }
                }
            }
        }
        
        return true;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

class Solution:
    def divide_teams(self, enemies: Dict[int, List[int]]) -> Union[List[List[int]], bool]:
        if not enemies:
            return [[], []]
        
        colors = {}
        students = set()
        
        # Get all students
        for student, enemy_list in enemies.items():
            students.add(student)
            for enemy in enemy_list:
                students.add(enemy)
        
        # Try to color using BFS
        for student in students:
            if student not in colors:
                if not self._bfs(student, enemies, colors):
                    return False
        
        # Build teams
        team1 = [student for student in students if colors[student] == 0]
        team2 = [student for student in students if colors[student] == 1]
        
        return [team1, team2]
    
    def _bfs(self, start: int, enemies: Dict[int, List[int]], 
             colors: Dict[int, int]) -> bool:
        from collections import deque
        
        queue = deque([start])
        colors[start] = 0
        
        while queue:
            curr = queue.popleft()
            
            if curr in enemies:
                for enemy in enemies[curr]:
                    if enemy not in colors:
                        colors[enemy] = 1 - colors[curr]
                        queue.append(enemy)
                    elif colors[enemy] == colors[curr]:
                        return False  # Conflict found
        
        return True

Complexity

  • ⏰ Time complexity: O(V + E), where V is the number of students and E is the number of enemy relationships. Each student and edge is processed exactly once
  • 🧺 Space complexity: O(V + E), for the color mapping, student set, and BFS queue which can hold up to V elements in the worst case