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

• DP형태로 점화식은 쉽게 세울수 있다. $ dp[i] = max_{j\in{P_i}}{dp[j]} + 1 $ 과 같은 형태이다. P_i는 S[i] 바로 앞에 올수있는 S[j]들의 인덱스들이다. 수식으로는 이런식..P_i = {j | j <i and |S[j] - S[i]| ≤ d or S[j] ≡ S[i]}
• 쉽게 생각하면, dp[i]를 계산할때 모든 j<i 에 대해서 S[j]가 S[i]앞에 올수 있는 조건을 만족하는지 확인하고, 만족하는 j들에 대해서 dp[j]의 최댓값을 구하면 된다. 이렇게 하면 O(n^2).
• O(n^2)은 너무 느리니까, 앞에 올수 있는 조건을 만족하는 S[j]들 중에서의 최댓값을 더 빠르게 구하는 방법을 생각하자.
• mod값이 같은 j들중 dp[j]의 최댓값을 구하는 것은 S[j]%K 가 같은 j들에 대해서 max(dp[j])를 저장해놓으면 바로 찾을수 있다. O(1)에 찾고, 업데이트도 O(1)에 가능
• S[j]가 범위 안에 들어오는 j들중 dp[j]의 최댓값을 구하는 것은 세그먼트 트리를 사용한다. S[j]를 인덱스로 두고, dp[j]를 밸류로 해서 최댓값 세그먼트 트리를 만들어두면, 최댓값을 찾는것과 업데이트를 둘다 O(logm)에 할수 있다. (m은 S[i]의 최댓값)
• n번의 쿼리를 처리해야 하므로 총 시간복잡도는 O(nlogm)이 된다.
• 다만, 이 문제에서는 n과 m이 최대 500,000으로 주어지는데, 파이썬으로는 통과시키기가 빡빡하다. 제너럴한 연산을 지원하는 teflib.segmenttree.SegmentTree을 사용했을때에는 통과가 안됐고, PyPy로 1000ms 정도가 걸렸다. 시간을 조금이라도 줄이기 위해서 열심히 최적화를 한 결과, 겨우 python으로도 5500ms 정도에 통과시킬수 있었다. 여기에서 가장 큰 최적화는, max함수를 호출하는 부분을 비교 연산자와 if문으로 바꿔 쓴 것이었는데 (라이브러리 부분과 메인부분 모두..), 특히 라이브러리부분을 max연산에 최적화한 것은 다른 문제에도 사용될 여지가 있어서, 따로 클래스를 만들어서 보관하기로 했다. max보다는 좀더 범용적인 min을 지원하도록 수정해서 teflib.segmenttree.MinSegmentTree 클래스로 만들었다. (max용으로 사용할때는 값들에 전부 -를 씌워서 넣으면 된다. heapq을 max heap으로 사용할때와 같은 방법.)

"""Solution code for "BOJ 15648. 추출하는 폴도 바리스타입니다".

- Problem link: https://www.acmicpc.net/problem/15648
- Solution link: http://www.teferi.net/ps/problems/boj/15648

Tags: [Segment tree] [DP]
"""

from teflib import segmenttree


def main():
    N, k, d = [int(x) for x in input().split()]
    S = [int(x) for x in input().split()]

    size = max(S) + 1
    segtree = segmenttree.MinSegmentTree(size)
    max_vals_for_mod = [0] * k
    for s_i in S:
        max_in_mod = max_vals_for_mod[(mod := s_i % k)]
        beg = 0 if s_i < d else s_i - d
        end = size if size < (p := s_i + d + 1) else p
        max_in_range = -segtree.query(beg, end)
        dp_i = (max_in_range if max_in_range > max_in_mod else max_in_mod) + 1

        max_vals_for_mod[mod] = dp_i
        segtree.set(s_i, -dp_i)

    print(max(max_vals_for_mod))


if __name__ == '__main__':
    main()

• Dependency: teflib.segmenttree.MinSegmentTree