SRM 404: KSubstring

(問題概要)
A[0] .. A[n - 1] (引数から線形合同法みたいなので構成)
に対して、s(i, k) = A[i] + ... + A[i + k - 1] とする。

s(i, k) - s(j, k)

(i + k - 1 < j)

を最小にするkとその値を求めよ。

(解法)
kを1ずつ増やしていった時に、s(i, k) (i = 0 .. n-1) の更新はO(n)でできる。
各jに対して目的の値を最小にするiは2分探索で求まる。
全体でO(n ^ 2 log n)。
2分探索をstd::lower_boundでやったらどうやらO(log n)にならないっぽく、std::set::lower_boundを使わなければいけない模様。
本番にはまらなくて良かった・・・。

#include <cstdlib>
#include <algorithm>
#include <cassert>
#include <climits>
typedef long long ll;

#define REP(i,n) for(int i=0; i<(int)(n); i++)

#include <queue>
#include <set>
#include <iostream>
#include <sstream>

using namespace std;

class KSubstring {
public:
  vector <int> maxSubstring(int A0, int X, int Y, int M, int n);
};

vector <int> KSubstring::maxSubstring(int A0, int X, int Y, int M, int n){
  vector<int> a(n);

  REP(i,n){
    a[i] = A0;
    A0 = ((ll)A0 * X + Y) % M;
  }

  vector<ll> s(a.begin(), a.end());

  ll val = 1ll << 60;
  int ans = 0;

  for(int k = 1; k <= n; k++){
    set<ll> memo;

    ll m = 1ll << 60;
    for(int i = k; i + k - 1 < n; i++){
      const ll v = s[i];

      memo.insert(s[i - k]);

      set<ll>::iterator it =
	memo.lower_bound(v);

      if(it != memo.end())
	m = min(m, std::abs(v - *it));
      if(it != memo.begin())
	m = min(m, std::abs(v - *(--it)));

      if(m <= val){
	val = m;
	ans = k;
      }
    }

    REP(j,n) if(j + k < n)
      s[j] += a[j + k];
  }

  vector<int> ret(2);
  ret[0] = ans;
  ret[1] = val;

  return ret;
}