1+ import java .util .LinkedList ;
2+ import java .util .Queue ;
3+ import java .util .TreeMap ;
4+
5+ public class Solution1825 {
6+ static class MKAverage {
7+ private final int m ;
8+ private final int k ;
9+ private final Queue <Integer > queue ;
10+ // 分别存储最小 k 个，中间，最大 k 个元素
11+ private final TreeMap <Integer , Integer > map1 , map2 , map3 ;
12+ // 三个 TreeMap 对应的 size
13+ private int sz1 , sz2 , sz3 ;
14+ // 容器剩余 m-2k 个元素的和
15+ private long sum ;
16+
17+ public MKAverage (int m , int k ) {
18+ this .m = m ;
19+ this .k = k ;
20+ queue = new LinkedList <>();
21+ map1 = new TreeMap <>();
22+ map2 = new TreeMap <>();
23+ map3 = new TreeMap <>();
24+ sz1 = 0 ;
25+ sz2 = 0 ;
26+ sz3 = 0 ;
27+ sum = 0L ;
28+ }
29+
30+ public void addElement (int num ) {
31+ queue .add (num );
32+ if (queue .size () <= m ) {
33+ map2add (num );
34+ if (queue .size () == m ) {
35+ while (sz1 < k ) {
36+ int min = map2 .firstKey ();
37+ map2del (min );
38+ map1add (min );
39+ }
40+ while (sz3 < k ) {
41+ int max = map2 .lastKey ();
42+ map2del (max );
43+ map3add (max );
44+ }
45+ }
46+ return ;
47+ }
48+
49+ // q 的元素数目等于 m+1
50+ if (num < map1 .lastKey ()) {
51+ map1add (num );
52+ int max = map1 .lastKey ();
53+ map1del (max );
54+ map2add (max );
55+ } else if (num > map3 .firstKey ()) {
56+ map3add (num );
57+ int min = map3 .firstKey ();
58+ map3del (min );
59+ map2add (min );
60+ } else {
61+ map2add (num );
62+ }
63+
64+ int rm = queue .remove ();
65+ if (map1 .containsKey (rm )) {
66+ map1del (rm );
67+ int min = map2 .firstKey ();
68+ map2del (min );
69+ map1add (min );
70+ } else if (map3 .containsKey (rm )) {
71+ map3del (rm );
72+ int max = map2 .lastKey ();
73+ map2del (max );
74+ map3add (max );
75+ } else {
76+ map2del (rm );
77+ }
78+ }
79+
80+ public int calculateMKAverage () {
81+ if (queue .size () < m ) {
82+ return -1 ;
83+ }
84+ return (int ) (sum / (m - k - k ));
85+ }
86+
87+ // 封装 "原子操作"
88+ private void map1del (int num ) {
89+ int val = map1 .get (num );
90+ if (val - 1 == 0 ) {
91+ map1 .remove (num );
92+ } else {
93+ map1 .put (num , val - 1 );
94+ }
95+ sz1 --;
96+ }
97+
98+ private void map2del (int num ) {
99+ int val = map2 .get (num );
100+ if (val - 1 == 0 ) {
101+ map2 .remove (num );
102+ } else {
103+ map2 .put (num , val - 1 );
104+ }
105+ sz2 --;
106+ sum -= num ;
107+ }
108+
109+ private void map3del (int num ) {
110+ int val = map3 .get (num );
111+ if (val - 1 == 0 ) {
112+ map3 .remove (num );
113+ } else {
114+ map3 .put (num , val - 1 );
115+ }
116+ sz3 --;
117+ }
118+
119+ private void map1add (int num ) {
120+ map1 .put (num , map1 .getOrDefault (num , 0 ) + 1 );
121+ sz1 ++;
122+ }
123+
124+ private void map2add (int num ) {
125+ map2 .put (num , map2 .getOrDefault (num , 0 ) + 1 );
126+ sz2 ++;
127+ sum += num ;
128+ }
129+
130+ private void map3add (int num ) {
131+ map3 .put (num , map3 .getOrDefault (num , 0 ) + 1 );
132+ sz3 ++;
133+ }
134+ }
135+ }
136+ /*
137+ 1825. 求出 MK 平均值
138+ https://leetcode.cn/problems/finding-mk-average/
139+
140+ 给你两个整数 m 和 k ，以及数据流形式的若干整数。你需要实现一个数据结构，计算这个数据流的 MK 平均值 。
141+ MK 平均值 按照如下步骤计算：
142+ 1. 如果数据流中的整数少于 m 个，MK 平均值 为 -1 ，否则将数据流中最后 m 个元素拷贝到一个独立的容器中。
143+ 2. 从这个容器中删除最小的 k 个数和最大的 k 个数。
144+ 3. 计算剩余元素的平均值，并 向下取整到最近的整数 。
145+ 请你实现 MKAverage 类：
146+ - MKAverage(int m, int k) 用一个空的数据流和两个整数 m 和 k 初始化 MKAverage 对象。
147+ - void addElement(int num) 往数据流中插入一个新的元素 num 。
148+ - int calculateMKAverage() 对当前的数据流计算并返回 MK 平均数 ，结果需 向下取整到最近的整数 。
149+ 提示：
150+ 3 <= m <= 10^5
151+ 1 <= k*2 < m
152+ 1 <= num <= 10^5
153+ addElement 与 calculateMKAverage 总操作次数不超过 10^5 次。
154+
155+ 队列 + 三个 TreeMap 模拟
156+ */
0 commit comments