SRM310 D1M-D2H FloatingMedian

TopCoder Statistics - Problem Statement

This problem statement is the exclusive and proprietary property of TopCoder, Inc. Any unauthorized use or reproduction of this information without the prior written consent of TopCoder, Inc. is strictly prohibited. (c)2010, TopCoder, Inc. All rights reserved.

乱数列を作り,連続する$K$個の部分数列の中央値の和を求める.愚直に求めると$O(N \times Klog(K))$となり間に合わない.重複なしの$k$番目を求めるような場合であれば平衡二分木かな?と思ったが,重複ありの,毎回求めるのは中央値だけで良い.中央値を$x$と置くと要素が$K$個の数列はsortした時に,indexが$x$より前のものと後のものに分けることが出来る.これを$down, up$とし,sortされている列とする.一個数列をずらすには左端の要素をpopし,右端$+1$の要素をpushすれば良い.
次のように場合分けした.

  • popする要素が$x$で,pushするのが$down$($x以下$) $\to$ pop, push後$down$の末尾を$x$にしてpop
  • popする要素が$x$で,pushするのが$up$($xより大きい$) $\to$ pop, push後$up$の先頭を$x$にしてpop
  • popする要素が$down$の中にあって, pushするのが$down$($x以下$) $\to$ pop, pushのみ
  • popする要素が$down$の中にあって, pushするのが$up$($xより大きい$) $\to$ pop, push後$x$を$down$にpush,$up$の先頭を$x$にしてpop
  • popする要素が$up$の中にあって, pushするのが$down$($x以下$) $\to$ pop, push後$x$を$up$にpush, $down$の末尾を$x$にしてpop
  • popする要素が$up$の中にあって, pushするのが$up$($xより大きい$) pop, pushのみ

$down, up$はmapで表現した.

Code

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
93
94
95
96
97
98
99
#include <iostream>
#include <sstream>
#include <vector>
#include <string>
#include <cstring>
#include <algorithm>
#include <map>
#include <set>
#include <queue>

#define REP(i,k,n) for(int i=k;i<n;i++)
#define rep(i,n) for(int i=0;i<n;i++)
#define each(it,v) for(__typeof((v).begin()) it=(v).begin();it!=(v).end();it++)
#define INF 1<<30
#define mp make_pair
#define MOD 65536

using namespace std;
typedef long long ll;
typedef pair<int, int> P;

class FloatingMedian {
  public:
  long long sumOfMedians(int seed, int mul, int add, int N, int K) {
      ll t = seed;
      vector<ll> v;
      rep(i, N) {
          v.push_back(t);
          t = (t * mul + add) % MOD;
      }

      int l = 0, r = K;

      vector<ll> res(r);
      rep(i, r) {
          res[i] = v[i];
      }
      sort(res.begin(), res.end());

      ll x = res[(K + 1) / 2 - 1];
      map<ll, int> down, up;
      rep(i, r) {
          if(i < (K + 1) / 2 - 1) down[res[i]]++;
          if(i > (K + 1) / 2 - 1) up[res[i]]++;
      }

      ll ans = x;
      int cnt = 0;
      while(r < N) {
          bool f1 = (down.count(v[l]) == 0);
          bool f2 = (up.count(v[l]) == 0);

          if(f1 && f2) {
              if(v[r] <= x) {
                  down[v[r]]++;
                  map<ll, int>::iterator it = down.end(); it--;
                  x = it->first;
                  down[x]--; if(down[x] == 0) down.erase(x);
              } else {
                  up[v[r]]++;
                  x = up.begin()->first;
                  up[x]--; if(up[x] == 0) up.erase(x);
              }
          } else if(!f1) {
              if(v[r] <= x) {
                  down[v[l]]--; if(down[v[l]] == 0) down.erase(v[l]);
                  down[v[r]]++;
              } else {
                  down[v[l]]--; if(down[v[l]] == 0) down.erase(v[l]);
                  down[x]++;
                  up[v[r]]++;

                  x = up.begin()->first;
                  up[x]--; if(up[x] == 0) up.erase(x);
              }
          } else if(!f2) {
              if(v[r] <= x) {
                  up[v[l]]--; if(up[v[l]] == 0) up.erase(v[l]);
                  up[x]++;
                  down[v[r]]++;

                  map<ll, int>::iterator it = down.end(); it--;
                  x = it->first;
                  down[x]--; if(down[x] == 0) down.erase(x);
              } else {
                  up[v[l]]--; if(up[v[l]] == 0) up.erase(v[l]);
                  up[v[r]]++;
              }
          }

          cnt++;
          ans += x;

          l++; r++;
      }
      
      return ans;
  }
};
Sep 9th, 2016