Algorithms-Library/Sliding Median.cpp at main · Alextz307/Algorithms-Library · GitHub

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
int median = -1;
multiset<int> low, big;
int64_t sum_low, sum_big;

void fix() {
  int m = (low.size() + big.size()) >> 1;
  if ((int)low.size() < m) {
    low.emplace(median);
    sum_low += median;
    median = *big.begin();
    big.erase(big.begin());
    sum_big -= median;
  }
  if ((int)low.size() > m) {
    big.emplace(median);
    sum_big += median;
    median = *--low.end();
    low.erase(--low.end());
    sum_low -= median;
  }
}

void add(int x) {
  if (median == -1) {
    median = x;
    return;
  }
  if (x < median) {
    low.emplace(x);
    sum_low += x;
  } else {
    big.emplace(x);
    sum_big += x;
  }
  fix();
}

void rem(int x) {
  if (x == median) {
    median = *big.begin();
    big.erase(big.begin());
    sum_big -= median;
  } else {
    if (x < median) {
      low.erase(low.find(x));
      sum_low -= x;
    } else {
      big.erase(big.find(x));
      sum_big -= x;
    }
  }
  fix();
}