Hatena::Grouptopcoder

tsubosakaの日記

2009-01-17

SRM 310 Div1 Medium

| 18:17

幅Kの窓を配列の0からN-K+1までシフトさせるときに各場所での窓中の要素の中央値を求めよという問題.

配列の要素が65536までしか取りえないのでbinary indexed treeを使ってsum(0,i) >= K / 2となる場所を二分探索を使って求める.

public class FloatingMedian {
  class FenwickTree{
    long x[];
    public FenwickTree(int n) {
      x = new long[n];
    }
    long sum(int i , int j){
      if(i == 0){
        long s = 0;
        for( ; j >= 0 ; j = (j & (j + 1)) - 1){
          s += x[j];
        }
        return s;
      }else{
        return sum(0 , j) - sum(0 , i - 1);
      }
    }
    void add(int k , long a){
      for(; k < x.length ; k |= k + 1){
        x[k] += a;
      }
    }
    long ithpoint(int index){
      int low = 0;
      int high = x.length;
      while(low < high){
        int mid = low + (high - low) / 2;
        if(sum(0 , mid) >= index)
          high = mid;
        else
          low = mid + 1;
      }
      return low;
    }
  }
  final static int MOD = 65536;
  public long sumOfMedians(int seed, int mul, int add, int N, int K) {
    long[] t = generate(seed, mul, add, N);
    int mindex = (K + 1) / 2;
    FenwickTree ft = new FenwickTree(MOD);
    for(int i = 0 ; i < K ; i++)
      ft.add((int)t[i], 1);
    long res = 0;
    for(int i = 0 ; i < N - K + 1 ; i++){
      long median = ft.ithpoint(mindex);
      res += median;
      if(i + K < t.length){
        ft.add((int)t[i], -1);
        ft.add((int)t[i + K], 1);        
      }
    }
    return res;
  }

  private long[] generate(int seed, int mul, int add, int N) {
    long t[] = new long[N];
    t[0] = seed;
    for(int i = 1 ; i < N ; i++){
      t[i] = t[i - 1] * mul + add;
      t[i] %= MOD;
    }
    return t;
  }
}