Finding MK Average LeetCode solution in Java

1
class MKAverage {
2
  class Node implements Comparable<Node> {
3
    int val;
4
    int time;
5

6
    Node(int val, int time) {
7
      this.val = val;
8
      this.time = time;
9
    }
10

11
    @Override
12
    public int compareTo(Node other) {
13
      return (this.val != other.val ? this.val - other.val : this.time - other.time);
14
    }
15
  }
16

17
  private TreeSet<Node> set = new TreeSet<>(); // natural order
18
  private Deque<Node> queue = new LinkedList<>();
19
  private Node kLeft;
20
  private Node kRight;
21

22
  private int m, k;
23

24
  private int time = 0;
25
  private int sum = 0;
26

27
  public MKAverage(int m, int k) {
28
    this.m = m;
29
    this.k = k;
30
  }
31

32
  public void addElement(int num) {
33
    Node node = new Node(num, time++);
34

35
    addNode(node);
36
    removeNode();
37

38
    if (time == m) init();
39
  }
40

41
  private void init() {
42
    int i = 0;
43
    for (Node node : set) {
44
      if (i < k - 1)
45
        ;
46
      else if (i == k - 1) kLeft = node;
47
      else if (i < m - k) sum += node.val;
48
      else if (i == m - k) {
49
        kRight = node;
50
        return;
51
      }
52

53
      i++;
54
    }
55
    return;
56
  }
57

58
  private void addNode(Node node) {
59
    queue.offerLast(node);
60
    set.add(node);
61

62
    if (queue.size() <= m) return;
63

64
    if (node.compareTo(kLeft) < 0) {
65
      sum += kLeft.val;
66
      kLeft = set.lower(kLeft);
67
    } else if (node.compareTo(kRight) > 0) {
68
      sum += kRight.val;
69
      kRight = set.higher(kRight);
70
    } else {
71
      sum += node.val;
72
    }
73
  }
74

75
  private void removeNode() {
76
    if (queue.size() <= m) return;
77

78
    Node node = queue.pollFirst();
79

80
    if (node.compareTo(kLeft) <= 0) {
81
      kLeft = set.higher(kLeft);
82
      sum -= kLeft.val;
83
    } else if (node.compareTo(kRight) >= 0) {
84
      kRight = set.lower(kRight);
85
      sum -= kRight.val;
86
    } else {
87
      sum -= node.val;
88
    }
89

90
    set.remove(node);
91
  }
92

93
  public int calculateMKAverage() {
94
    return (queue.size() < m ? -1 : sum / (m - 2 * k));
95
  }
96
}