Skip to content

1825. Finding MK Average

Approach 1: Lots variables + while

  • Time: $O(n\log n)$
  • Space: $O(n)$
 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
class MKAverage {
 public:
  MKAverage(int m, int k) : m(m), k(k) {}

  void addElement(int num) {
    q.push(num);
    add(mid, num);
    midSum += num;

    if (q.size() > m) {
      const int removed = q.front();
      q.pop();
      if (top.contains(removed)) {
        remove(top, removed);
        --topSize;
      } else if (mid.contains(removed)) {
        remove(mid, removed);
        midSum -= removed;
      } else {
        remove(bot, removed);
        --botSize;
      }
    }

    // Move item(s) from `mid` to `top` to fill k slots.
    while (!mid.empty() && topSize < k) {
      midSum -= mid.rbegin()->first;
      add(top, remove(mid, mid.rbegin()->first));
      ++topSize;
    }

    // Rebalance `mid` and `top`.
    while (!mid.empty() && mid.rbegin()->first > top.begin()->first) {
      midSum -= mid.rbegin()->first;
      midSum += top.begin()->first;
      add(top, remove(mid, mid.rbegin()->first));
      add(mid, remove(top, top.begin()->first));
    }

    // Move item(s) from `mid` to `bot` to fill k slots.
    while (!mid.empty() && botSize < k) {
      midSum -= mid.begin()->first;
      add(bot, remove(mid, mid.begin()->first));
      ++botSize;
    }

    // Move item(s) from `mid` to `bot` to fill k slots.
    while (!mid.empty() && mid.begin()->first < bot.rbegin()->first) {
      midSum -= mid.begin()->first;
      midSum += bot.rbegin()->first;
      add(bot, remove(mid, mid.begin()->first));
      add(mid, remove(bot, bot.rbegin()->first));
    }
  }

  int calculateMKAverage() {
    return q.size() == m ? midSum / (m - 2 * k) : -1;
  }

 private:
  const int m;
  const int k;
  queue<int> q;
  map<int, int> top;
  map<int, int> mid;
  map<int, int> bot;
  int topSize = 0;
  int botSize = 0;
  long midSum = 0;

  void add(map<int, int>& map, int num) {
    ++map[num];
  }

  int remove(map<int, int>& map, int num) {
    if (--map[num] == 0)
      map.erase(num);
    return num;
  }
};

 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
class MKAverage {
  public MKAverage(int m, int k) {
    this.m = m;
    this.k = k;
  }

  public void addElement(int num) {
    q.offer(num);
    add(mid, num);
    midSum += num;

    if (q.size() > m) {
      final int removed = q.poll();
      if (top.containsKey(removed)) {
        remove(top, removed);
        --topSize;
      } else if (mid.containsKey(removed)) {
        remove(mid, removed);
        midSum -= removed;
      } else {
        remove(bot, removed);
        --botSize;
      }
    }

    // Move item(s) from `mid` to `top` to fill k slots.
    while (!mid.isEmpty() && topSize < k) {
      midSum -= mid.lastKey();
      add(top, remove(mid, mid.lastKey()));
      ++topSize;
    }

    // Rebalance `mid` and `top`.
    while (!mid.isEmpty() && mid.lastKey() > top.firstKey()) {
      midSum -= mid.lastKey();
      midSum += top.firstKey();
      add(top, remove(mid, mid.lastKey()));
      add(mid, remove(top, top.firstKey()));
    }

    // Move item(s) from `mid` to `bot` to fill k slots.
    while (!mid.isEmpty() && botSize < k) {
      midSum -= mid.firstKey();
      add(bot, remove(mid, mid.firstKey()));
      ++botSize;
    }

    // Rebalance mid and bot
    while (!mid.isEmpty() && mid.firstKey() < bot.lastKey()) {
      midSum -= mid.firstKey();
      midSum += bot.lastKey();
      add(bot, remove(mid, mid.firstKey()));
      add(mid, remove(bot, bot.lastKey()));
    }
  }

  public int calculateMKAverage() {
    return q.size() == m ? (int) (midSum / (m - 2 * k)) : -1;
  }

  private final int m;
  private final int k;
  private Queue<Integer> q = new ArrayDeque<>();
  private TreeMap<Integer, Integer> top = new TreeMap<>();
  private TreeMap<Integer, Integer> mid = new TreeMap<>();
  private TreeMap<Integer, Integer> bot = new TreeMap<>();
  private int topSize = 0;
  private int botSize = 0;
  private long midSum = 0;

  private void add(TreeMap<Integer, Integer> map, int num) {
    map.merge(num, 1, Integer::sum);
  }

  private int remove(TreeMap<Integer, Integer> map, int num) {
    map.put(num, map.get(num) - 1);
    if (map.get(num) == 0)
      map.remove(num);
    return num;
  }
}

Approach 2: Wrap with MyMap + if

  • Time: $O(n\log n)$
  • Space: $O(n)$
 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
struct MyMap {
  map<int, int> map;
  int size = 0;
  long sum = 0;
};

class MKAverage {
 public:
  MKAverage(int m, int k) : m(m), k(k), kMidSize(m - 2 * k) {}

  void addElement(int num) {
    q.push(num);
    add(num);

    if (q.size() > m) {
      const int removed = q.front();
      q.pop();
      remove(removed);
    }
  }

  int calculateMKAverage() {
    return q.size() == m ? mid.sum / kMidSize : -1;
  }

 private:
  const int m;
  const int k;
  const int kMidSize;
  queue<int> q;
  MyMap top;
  MyMap mid;
  MyMap bot;

  void add(int num) {
    add(bot, num);
    if (bot.size > k)
      add(mid, remove(bot, bot.map.rbegin()->first));
    if (mid.size > kMidSize)
      add(top, remove(mid, mid.map.rbegin()->first));
  }

  void remove(int num) {
    if (bot.map.contains(num))
      remove(bot, num);
    else if (mid.map.contains(num))
      remove(mid, num);
    else
      remove(top, num);

    if (bot.size < k)
      add(bot, remove(mid, mid.map.begin()->first));
    if (mid.size < kMidSize)
      add(mid, remove(top, top.map.begin()->first));
  }

  void add(MyMap& m, int num) {
    ++m.map[num];
    ++m.size;
    m.sum += num;
  }

  int remove(MyMap& m, int num) {
    if (--m.map[num] == 0)
      m.map.erase(num);
    --m.size;
    m.sum -= num;
    return num;
  }
};

 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
class MyMap {
  public TreeMap<Integer, Integer> map = new TreeMap<>();
  public int size = 0;
  public long sum = 0;
}

class MKAverage {
  public MKAverage(int m, int k) {
    this.m = m;
    this.k = k;
    this.kMidSize = m - 2 * k;
  }

  public void addElement(int num) {
    q.offer(num);
    add(num);

    if (q.size() > m)
      remove(q.poll());
  }

  public int calculateMKAverage() {
    return q.size() == m ? (int) (mid.sum / kMidSize) : -1;
  }

  private final int m;
  private final int k;
  private final int kMidSize;
  private Queue<Integer> q = new ArrayDeque<>();
  private MyMap bot = new MyMap();
  private MyMap mid = new MyMap();
  private MyMap top = new MyMap();

  private void add(int num) {
    add(bot, num);
    if (bot.size > k)
      add(mid, remove(bot, bot.map.lastKey()));
    if (mid.size > kMidSize)
      add(top, remove(mid, mid.map.lastKey()));
  }

  private void remove(int num) {
    if (bot.map.containsKey(num))
      remove(bot, num);
    else if (mid.map.containsKey(num))
      remove(mid, num);
    else
      remove(top, num);

    if (bot.size < k)
      add(bot, remove(mid, mid.map.firstKey()));
    if (mid.size < kMidSize)
      add(mid, remove(top, top.map.firstKey()));
  }

  private void add(MyMap m, int num) {
    m.map.merge(num, 1, Integer::sum);
    ++m.size;
    m.sum += num;
  }

  private int remove(MyMap m, int num) {
    m.map.merge(num, -1, Integer::sum);
    if (m.map.get(num) == 0)
      m.map.remove(num);
    --m.size;
    m.sum -= num;
    return num;
  }
}