BOJ 15648 - 추출하는 폴도 바리스타입니다

BOJ 15648 - 추출하는 폴도 바리스타입니다

naive하게 $dp$ 식을 세우면

$dp_{i}$는 $1+Max_{j=i+1}^{n-1} \{dp_{j} ~\text{if}~ a_{i}\equiv a_{j}~(mod~k)\}$ 이나

$1+Max_{j=i+1}^{n-1} \{dp_{j} ~\text{if}~ a_{j}-d \le a_{i}\le a_{j}+d \}$ 둘 중에 큰 값이여야 하고 구하는데 $O(n^2)$ 이 걸린다.

따라서 이걸 세그먼트 트리로 최적화하여 $O(n\log n)$에 문제를 해결하도록 한다.

두 세그먼트 트리를 준비하고 $dp$ 값에 대한 지금까지 나온 범위중에 $max$ 값을 찾도록 한다.

void solve() {  
   int n, k, d;  
   cin >> n >> k >> d;  
   vi a(n);  
   fv(a);  
  
   seg_tree<int> dp_sum(6e5), dp_mod(k + 5);  
  
   int ans = 0;  
   for (int i = n - 1; i >= 0; i--) {  
      int mod = a[i] % k;  
      int cur = dp_mod.query(mod, mod).max;  
      cur = max(cur, dp_sum.query(max(1, a[i] - d), min(a[i] + d, int(5e5))).max);  
  
      cur++;  
      ans = max(ans, cur);  
  
      dp_mod.update(mod, cur);  
      dp_sum.update(a[i], cur);  
   }  
   cout << ans;  
}