TopCoder SRM 310 Div1 Medium FloatingMedian
問題
数列t[i]が与えられる。
長さKの連続する部分列の中央値の、全ての合計を求めよ。
ただし、長さnの数列の中央値とは、nを小さい順に並び替えたときにn/2番目の要素を指すものとする。
制約条件
t[i]の長さ≦250000
K≦5000
0≦t[i]≦65535
方針
二分探索木にn番目の要素をO(log size)で求める機能をつけたものがspaghetti sourceにあったので使った。
単に要素を追加していって、k/2番目の要素を答えに足して、要素を消していくだけ。
想定解は、t[i]≦65536であることを使って、
binary indexed treeを使って集合を表現し、n番目の要素をO(log n^2)（→実はO(logn)だった。追記）で求めるというもの。
ソースコード
template <class T> struct avl_tree { struct node { T key; int size, height; node *child[2]; node(const T &key) : key(key), size(1), height(1) { child[0] = child[1] = 0; } } *root; typedef node *pointer; avl_tree() { root = NULL; } pointer find(const T &key) { return find(root, key); } node *find(node *t, const T &key) { if (t == NULL) return NULL; if (key == t->key) return t; else if (key < t->key) return find(t->child[0], key); else return find(t->child[1], key); } void insert(const T &key) { root = insert(root, new node(key)); } node *insert(node *t, node *x) { if (t == NULL) return x; if (x->key <= t->key) t->child[0] = insert(t->child[0], x); else t->child[1] = insert(t->child[1], x); t->size += 1; return balance(t); } void erase(const T &key) { root = erase(root, key); } node *erase(node *t, const T &x) { if (t == NULL) return NULL; if (x == t->key) { return move_down(t->child[0], t->child[1]); } else { if (x < t->key) t->child[0] = erase(t->child[0], x); else t->child[1] = erase(t->child[1], x); t->size -= 1; return balance(t); } } node *move_down(node *t, node *rhs) { if (t == NULL) return rhs; t->child[1] = move_down(t->child[1], rhs); return balance(t); } #define sz(t) (t ? t->size : 0) #define ht(t) (t ? t->height : 0) node *rotate(node *t, int l, int r) { node *s = t->child[r]; t->child[r] = s->child[l]; s->child[l] = balance(t); if (t) t->size = sz(t->child[0]) + sz(t->child[1]) + 1; if (s) s->size = sz(s->child[0]) + sz(s->child[1]) + 1; return balance(s); } node *balance(node *t) { for (int i = 0; i < 2; ++i) { if (ht(t->child[!i]) - ht(t->child[i]) < -1) { if (ht(t->child[i]->child[!i]) - ht(t->child[i]->child[i]) > 0) t->child[i] = rotate(t->child[i], i, !i); return rotate(t, !i, i); } } if (t) t->height = max(ht(t->child[0]), ht(t->child[1])) + 1; if (t) t->size = sz(t->child[0]) + sz(t->child[1]) + 1; return t; } pointer rank(int k) const { return rank(root, k); } pointer rank(node *t, int k) const { if (!t) return NULL; int m = sz(t->child[0]); if (k < m) return rank(t->child[0], k); if (k == m) return t; if (k > m) return rank(t->child[1], k - m - 1); } }; int t[300000]; class FloatingMedian { public: long long sumOfMedians(int seed, int mul, int add, int N, int K) { t[0]=seed; rep(i,N)t[i+1]=(t[i]*(ll)mul+add)%65536; avl_tree<int> tree; ll ans=0; rep(i,N){ tree.insert(t[i]); if(i>=K-1){ ans+=tree.rank((K-1)/2)->key; tree.erase(t[i-K+1]); } } return ans; } };
追記
実はbitのk番目の値を求めるというのはO(logn)で出来た。
二分探索のときに、kからbit[mid]を引いていくような二分探索をすればよかった。
const int MX = 65536; int t[300000]; ll bit[MX + 1]; void add(int i, int x){ for(; i <= MX; i += i & -i) bit[i] += x; } inline ll sum(int i){ ll res = 0; for(; i; i -= i & -i) res += bit[i]; return res; } class FloatingMedian { public: long long sumOfMedians(int seed, int mul, int a, int N, int K) { rep(i, N){ t[i] = seed; seed = ((ll)seed * mul + a) % 65536; } memset(bit, 0, sizeof(bit)); ll res = 0; rep(i, N){ add(t[i] + 1, 1); if(i >= K - 1){ int lo = 0, hi = MX, mid, k = (K + 1) / 2; while(lo + 1 < hi){ mid = (lo + hi) / 2; if(bit[mid] >= k) hi = mid; else{ lo = mid; k -= bit[mid]; } } res += hi - 1; add(t[i - K + 1] + 1, -1); } } return res; } };
