Skip to content

1825. Finding MK Average

Description

You are given two integers, m and k, and a stream of integers. You are tasked to implement a data structure that calculates the MKAverage for the stream.

The MKAverage can be calculated using these steps:

  1. If the number of the elements in the stream is less than m you should consider the MKAverage to be -1. Otherwise, copy the last m elements of the stream to a separate container.
  2. Remove the smallest k elements and the largest k elements from the container.
  3. Calculate the average value for the rest of the elements rounded down to the nearest integer.

Implement the MKAverage class:

  • MKAverage(int m, int k) Initializes the MKAverage object with an empty stream and the two integers m and k.
  • void addElement(int num) Inserts a new element num into the stream.
  • int calculateMKAverage() Calculates and returns the MKAverage for the current stream rounded down to the nearest integer.

 

Example 1:

Input
["MKAverage", "addElement", "addElement", "calculateMKAverage", "addElement", "calculateMKAverage", "addElement", "addElement", "addElement", "calculateMKAverage"]
[[3, 1], [3], [1], [], [10], [], [5], [5], [5], []]
Output
[null, null, null, -1, null, 3, null, null, null, 5]

Explanation
MKAverage obj = new MKAverage(3, 1); 
obj.addElement(3);        // current elements are [3]
obj.addElement(1);        // current elements are [3,1]
obj.calculateMKAverage(); // return -1, because m = 3 and only 2 elements exist.
obj.addElement(10);       // current elements are [3,1,10]
obj.calculateMKAverage(); // The last 3 elements are [3,1,10].
                          // After removing smallest and largest 1 element the container will be [3].
                          // The average of [3] equals 3/1 = 3, return 3
obj.addElement(5);        // current elements are [3,1,10,5]
obj.addElement(5);        // current elements are [3,1,10,5,5]
obj.addElement(5);        // current elements are [3,1,10,5,5,5]
obj.calculateMKAverage(); // The last 3 elements are [5,5,5].
                          // After removing smallest and largest 1 element the container will be [5].
                          // The average of [5] equals 5/1 = 5, return 5

 

Constraints:

  • 3 <= m <= 105
  • 1 < k*2 < m
  • 1 <= num <= 105
  • At most 105 calls will be made to addElement and calculateMKAverage.

Solutions

Solution 1: Ordered Set + Queue

Thinking

We need the average of an \(m\)-length window after dropping the smallest and largest \(k\) values, with frequent insertions. Sorting the window every time is \(O(m\log m)\) and too slow.

A queue stores insertion order. Three ordered multisets hold the lowest \(k\), the middle segment, and the highest \(k\), together with the middle sum \(s\). After each insert or pop we move overflow elements so both ends have size \(k\). The average is \(s/(m-2k)\).

We can maintain the following data structures or variables:

  • A queue \(q\) of length \(m\), where the head of the queue is the earliest added element, and the tail of the queue is the most recently added element;
  • Three ordered sets, namely \(lo\), \(mid\), \(hi\), where \(lo\) and \(hi\) store the smallest \(k\) elements and the largest \(k\) elements respectively, and \(mid\) stores the remaining elements;
  • A variable \(s\), maintaining the sum of all elements in \(mid\);
  • Some programming languages (such as Java, Go) additionally maintain two variables \(size1\) and \(size3\), representing the number of elements in \(lo\) and \(hi\) respectively.

When calling the \(addElement(num)\) function, perform the following operations in order:

  1. If \(lo\) is empty, or \(num \leq max(lo)\), then add \(num\) to \(lo\); otherwise if \(hi\) is empty, or \(num \geq min(hi)\), then add \(num\) to \(hi\); otherwise add \(num\) to \(mid\), and add the value of \(num\) to \(s\).
  2. Next, add \(num\) to the queue \(q\). If the length of the queue \(q\) is greater than \(m\) at this time, remove the head element \(x\) from the queue \(q\), then choose one of \(lo\), \(mid\) or \(hi\) that contains \(x\), and remove \(x\) from this set. If the set is \(mid\), subtract the value of \(x\) from \(s\).
  3. If the length of \(lo\) is greater than \(k\), then repeatedly remove the maximum value \(max(lo)\) from \(lo\), add \(max(lo)\) to \(mid\), and add the value of \(max(lo)\) to \(s\).
  4. If the length of \(hi\) is greater than \(k\), then repeatedly remove the minimum value \(min(hi)\) from \(hi\), add \(min(hi)\) to \(mid\), and add the value of \(min(hi)\) to \(s\).
  5. If the length of \(lo\) is less than \(k\) and \(mid\) is not empty, then repeatedly remove the minimum value \(min(mid)\) from \(mid\), add \(min(mid)\) to \(lo\), and subtract the value of \(min(mid)\) from \(s\).
  6. If the length of \(hi\) is less than \(k\) and \(mid\) is not empty, then repeatedly remove the maximum value \(max(mid)\) from \(mid\), add \(max(mid)\) to \(hi\), and subtract the value of \(max(mid)\) from \(s\).

When calling the \(calculateMKAverage()\) function, if the length of \(q\) is less than \(m\), return \(-1\), otherwise return \(\frac{s}{m - 2k}\).

In terms of time complexity, each call to the \(addElement(num)\) function has a time complexity of \(O(\log m)\), and each call to the \(calculateMKAverage()\) function has a time complexity of \(O(1)\). The space complexity is \(O(m)\).

 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
class MKAverage:
    def __init__(self, m: int, k: int):
        self.m = m
        self.k = k
        self.s = 0
        self.q = deque()
        self.lo = SortedList()
        self.mid = SortedList()
        self.hi = SortedList()

    def addElement(self, num: int) -> None:
        if not self.lo or num <= self.lo[-1]:
            self.lo.add(num)
        elif not self.hi or num >= self.hi[0]:
            self.hi.add(num)
        else:
            self.mid.add(num)
            self.s += num
        self.q.append(num)
        if len(self.q) > self.m:
            x = self.q.popleft()
            if x in self.lo:
                self.lo.remove(x)
            elif x in self.hi:
                self.hi.remove(x)
            else:
                self.mid.remove(x)
                self.s -= x
        while len(self.lo) > self.k:
            x = self.lo.pop()
            self.mid.add(x)
            self.s += x
        while len(self.hi) > self.k:
            x = self.hi.pop(0)
            self.mid.add(x)
            self.s += x
        while len(self.lo) < self.k and self.mid:
            x = self.mid.pop(0)
            self.lo.add(x)
            self.s -= x
        while len(self.hi) < self.k and self.mid:
            x = self.mid.pop()
            self.hi.add(x)
            self.s -= x

    def calculateMKAverage(self) -> int:
        return -1 if len(self.q) < self.m else self.s // (self.m - 2 * self.k)


# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# obj.addElement(num)
# param_2 = obj.calculateMKAverage()

 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
class MKAverage {

    private int m, k;
    private long s;
    private int size1, size3;
    private Deque<Integer> q = new ArrayDeque<>();
    private TreeMap<Integer, Integer> lo = new TreeMap<>();
    private TreeMap<Integer, Integer> mid = new TreeMap<>();
    private TreeMap<Integer, Integer> hi = new TreeMap<>();

    public MKAverage(int m, int k) {
        this.m = m;
        this.k = k;
    }

    public void addElement(int num) {
        if (lo.isEmpty() || num <= lo.lastKey()) {
            lo.merge(num, 1, Integer::sum);
            ++size1;
        } else if (hi.isEmpty() || num >= hi.firstKey()) {
            hi.merge(num, 1, Integer::sum);
            ++size3;
        } else {
            mid.merge(num, 1, Integer::sum);
            s += num;
        }
        q.offer(num);
        if (q.size() > m) {
            int x = q.poll();
            if (lo.containsKey(x)) {
                if (lo.merge(x, -1, Integer::sum) == 0) {
                    lo.remove(x);
                }
                --size1;
            } else if (hi.containsKey(x)) {
                if (hi.merge(x, -1, Integer::sum) == 0) {
                    hi.remove(x);
                }
                --size3;
            } else {
                if (mid.merge(x, -1, Integer::sum) == 0) {
                    mid.remove(x);
                }
                s -= x;
            }
        }
        for (; size1 > k; --size1) {
            int x = lo.lastKey();
            if (lo.merge(x, -1, Integer::sum) == 0) {
                lo.remove(x);
            }
            mid.merge(x, 1, Integer::sum);
            s += x;
        }
        for (; size3 > k; --size3) {
            int x = hi.firstKey();
            if (hi.merge(x, -1, Integer::sum) == 0) {
                hi.remove(x);
            }
            mid.merge(x, 1, Integer::sum);
            s += x;
        }
        for (; size1 < k && !mid.isEmpty(); ++size1) {
            int x = mid.firstKey();
            if (mid.merge(x, -1, Integer::sum) == 0) {
                mid.remove(x);
            }
            s -= x;
            lo.merge(x, 1, Integer::sum);
        }
        for (; size3 < k && !mid.isEmpty(); ++size3) {
            int x = mid.lastKey();
            if (mid.merge(x, -1, Integer::sum) == 0) {
                mid.remove(x);
            }
            s -= x;
            hi.merge(x, 1, Integer::sum);
        }
    }

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

/**
 * Your MKAverage object will be instantiated and called as such:
 * MKAverage obj = new MKAverage(m, k);
 * obj.addElement(num);
 * int param_2 = obj.calculateMKAverage();
 */

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

    void addElement(int num) {
        if (lo.empty() || num <= *lo.rbegin()) {
            lo.insert(num);
        } else if (hi.empty() || num >= *hi.begin()) {
            hi.insert(num);
        } else {
            mid.insert(num);
            s += num;
        }

        q.push(num);
        if (q.size() > m) {
            int x = q.front();
            q.pop();
            if (lo.find(x) != lo.end()) {
                lo.erase(lo.find(x));
            } else if (hi.find(x) != hi.end()) {
                hi.erase(hi.find(x));
            } else {
                mid.erase(mid.find(x));
                s -= x;
            }
        }
        while (lo.size() > k) {
            int x = *lo.rbegin();
            lo.erase(prev(lo.end()));
            mid.insert(x);
            s += x;
        }
        while (hi.size() > k) {
            int x = *hi.begin();
            hi.erase(hi.begin());
            mid.insert(x);
            s += x;
        }
        while (lo.size() < k && mid.size()) {
            int x = *mid.begin();
            mid.erase(mid.begin());
            s -= x;
            lo.insert(x);
        }
        while (hi.size() < k && mid.size()) {
            int x = *mid.rbegin();
            mid.erase(prev(mid.end()));
            s -= x;
            hi.insert(x);
        }
    }

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

private:
    int m, k;
    long long s = 0;
    queue<int> q;
    multiset<int> lo, mid, hi;
};

/**
 * Your MKAverage object will be instantiated and called as such:
 * MKAverage* obj = new MKAverage(m, k);
 * obj->addElement(num);
 * int param_2 = obj->calculateMKAverage();
 */

 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
type MKAverage struct {
    lo, mid, hi  *redblacktree.Tree
    q            []int
    m, k, s      int
    size1, size3 int
}

func Constructor(m int, k int) MKAverage {
    lo := redblacktree.NewWithIntComparator()
    mid := redblacktree.NewWithIntComparator()
    hi := redblacktree.NewWithIntComparator()
    return MKAverage{lo, mid, hi, []int{}, m, k, 0, 0, 0}
}

func (this *MKAverage) AddElement(num int) {
    merge := func(rbt *redblacktree.Tree, key, value int) {
        if v, ok := rbt.Get(key); ok {
            nxt := v.(int) + value
            if nxt == 0 {
                rbt.Remove(key)
            } else {
                rbt.Put(key, nxt)
            }
        } else {
            rbt.Put(key, value)
        }
    }

    if this.lo.Empty() || num <= this.lo.Right().Key.(int) {
        merge(this.lo, num, 1)
        this.size1++
    } else if this.hi.Empty() || num >= this.hi.Left().Key.(int) {
        merge(this.hi, num, 1)
        this.size3++
    } else {
        merge(this.mid, num, 1)
        this.s += num
    }
    this.q = append(this.q, num)
    if len(this.q) > this.m {
        x := this.q[0]
        this.q = this.q[1:]
        if _, ok := this.lo.Get(x); ok {
            merge(this.lo, x, -1)
            this.size1--
        } else if _, ok := this.hi.Get(x); ok {
            merge(this.hi, x, -1)
            this.size3--
        } else {
            merge(this.mid, x, -1)
            this.s -= x
        }
    }
    for ; this.size1 > this.k; this.size1-- {
        x := this.lo.Right().Key.(int)
        merge(this.lo, x, -1)
        merge(this.mid, x, 1)
        this.s += x
    }
    for ; this.size3 > this.k; this.size3-- {
        x := this.hi.Left().Key.(int)
        merge(this.hi, x, -1)
        merge(this.mid, x, 1)
        this.s += x
    }
    for ; this.size1 < this.k && !this.mid.Empty(); this.size1++ {
        x := this.mid.Left().Key.(int)
        merge(this.mid, x, -1)
        this.s -= x
        merge(this.lo, x, 1)
    }
    for ; this.size3 < this.k && !this.mid.Empty(); this.size3++ {
        x := this.mid.Right().Key.(int)
        merge(this.mid, x, -1)
        this.s -= x
        merge(this.hi, x, 1)
    }
}

func (this *MKAverage) CalculateMKAverage() int {
    if len(this.q) < this.m {
        return -1
    }
    return this.s / (this.m - 2*this.k)
}

/**
 * Your MKAverage object will be instantiated and called as such:
 * obj := Constructor(m, k);
 * obj.AddElement(num);
 * param_2 := obj.CalculateMKAverage();
 */

Solution 2: Single Ordered Set + Sliding Window

Thinking

Solution 1 juggles three trees and their size invariants. A single ordered list of the window is enough: build it when the window first fills, then adjust \(s\) by the rank of the inserted and deleted values. The structure is simpler and each update remains logarithmic.

Use a queue for insertion order and one ordered set for the current window of length \(m\). When the window first becomes full, build the set and sum the middle segment after dropping the smallest and largest \(k\) elements. Later insertions and deletions update that middle sum from the element's rank in the set.

Unlike method 1, which partitions the window into three sets, this approach keeps a single ordered sequence.

Each addElement call takes \(O(\log m)\) time, and each calculateMKAverage call takes \(O(1)\) time. The space complexity is \(O(m)\).

 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
class MKAverage:
    def __init__(self, m: int, k: int):
        self.m = m
        self.k = k
        self.sl = SortedList()
        self.q = deque()
        self.s = 0

    def addElement(self, num: int) -> None:
        self.q.append(num)
        if len(self.q) == self.m:
            self.sl = SortedList(self.q)
            self.s = sum(self.sl[self.k : -self.k])
        elif len(self.q) > self.m:
            i = self.sl.bisect_left(num)
            if i < self.k:
                self.s += self.sl[self.k - 1]
            elif self.k <= i <= self.m - self.k:
                self.s += num
            else:
                self.s += self.sl[self.m - self.k]
            self.sl.add(num)

            x = self.q.popleft()
            i = self.sl.bisect_left(x)
            if i < self.k:
                self.s -= self.sl[self.k]
            elif self.k <= i <= self.m - self.k:
                self.s -= x
            else:
                self.s -= self.sl[self.m - self.k]
            self.sl.remove(x)

    def calculateMKAverage(self) -> int:
        return -1 if len(self.sl) < self.m else self.s // (self.m - self.k * 2)


# Your MKAverage object will be instantiated and called as such:
# obj = MKAverage(m, k)
# obj.addElement(num)
# param_2 = obj.calculateMKAverage()

Comments