Find All People With Secret LeetCode solution in Java

1
class Solution {
2
  public List<Integer> findAllPeople(int n, int[][] meetings, int firstPerson) {
3

4
    // create <time, index> map
5
    Map<Integer, List<Integer>> timeToIndexes = new TreeMap<>();
6
    int m = meetings.length;
7
    for (int i = 0; i < m; i++) {
8
      timeToIndexes.putIfAbsent(meetings[i][2], new ArrayList<>());
9
      timeToIndexes.get(meetings[i][2]).add(i);
10
    }
11

12
    UF uf = new UF(n);
13
    // base
14
    uf.union(0, firstPerson);
15

16
    // for every time we have a pool of people that talk to each other
17
    // if someone knows a secret proir to this meeting - all pool will too
18
    // if not - reset unions from this pool
19
    for (int time : timeToIndexes.keySet()) {
20
      Set<Integer> pool = new HashSet<>();
21

22
      for (int ind : timeToIndexes.get(time)) {
23
        int[] currentMeeting = meetings[ind];
24
        uf.union(currentMeeting[0], currentMeeting[1]);
25
        pool.add(currentMeeting[0]);
26
        pool.add(currentMeeting[1]);
27
      }
28

29
      // meeting that took place now should't affect future
30
      // meetings if people don't know the secret
31
      for (int i : pool) if (!uf.connected(0, i)) uf.reset(i);
32
    }
33

34
    // if the person is conneted to 0 - they know a secret
35
    List<Integer> ans = new ArrayList<>();
36
    for (int i = 0; i < n; i++) if (uf.connected(i, 0)) ans.add(i);
37
    return ans;
38
  }
39

40
  // regular union find
41
  private static class UF {
42
    int[] parent, rank;
43

44
    public UF(int n) {
45
      parent = new int[n];
46
      rank = new int[n];
47
      for (int i = 0; i < n; i++) parent[i] = i;
48
    }
49

50
    public void union(int p, int q) {
51
      int rootP = find(p);
52
      int rootQ = find(q);
53

54
      if (rootP == rootQ) return;
55

56
      if (rank[rootP] < rank[rootQ]) {
57
        parent[rootP] = rootQ;
58
      } else {
59
        parent[rootQ] = rootP;
60
        rank[rootP]++;
61
      }
62
    }
63

64
    public int find(int p) {
65
      while (parent[p] != p) {
66
        p = parent[parent[p]];
67
      }
68
      return p;
69
    }
70

71
    public boolean connected(int p, int q) {
72
      return find(p) == find(q);
73
    }
74

75
    public void reset(int p) {
76
      parent[p] = p;
77
      rank[p] = 0;
78
    }
79
  }
80
}