Problem p02851 (python)

Fixed submission:
def read():
    N, K = list(map(int, input().strip().split()))
    A = list(map(int, input().strip().split()))
    return N, K, A

def solve(N, K, A):
    if K == 1:
        return 0
    
    S = [0 for _ in range(N+1)]
    for i in range(N):
        S[i+1] = S[i] + A[i]
    
    count_sum = 0
    count = dict()
    for j in range(1, N+1):
        j_mod_k = (S[j] - j) % K  # 0 <= jmodk < K
        if j == 1:
            for i in range(max(0, j-K+1), j):
                # (S[i] - i) % K の値 i_mod_k は、次のループでもほぼ重複する
                # i = max(0, j-K+1) の時の i_mod_k は、i+1で破棄する
                # i = j の時の i_mod_k は、i+1で生成する
                i_mod_k = (S[i] - i) % K
                if i_mod_k in count:
                    count[i_mod_k] += 1
                else:
                    count[i_mod_k] = 1
        else:
            if max(0, j-K+1) != max(0, (j-1)-K+1):
                i = max(0, (j-1)-K+1)
                i_mod_k = (S[i] - i) % K
                if i_mod_k in count:
                    if count[i_mod_k] == 1:
                        count.pop(i_mod_k)
                    else:
                        count[i_mod_k] -= 1
            i = j-1
            i_mod_k = (S[i] - i) % K
            if i_mod_k in count:
                count[i_mod_k] += 1
            else:
                count[i_mod_k] = 1
        if j_mod_k in count:
            count_sum += count[j_mod_k]
    return count_sum

if __name__ == '__main__':
    inputs = read()
    print('%d' % solve(*inputs))

Buggy submission (verdict: Runtime Error):
def read():
    N, K = list(map(int, input().strip().split()))
    A = list(map(int, input().strip().split()))
    return N, K, A

def solve(N, K, A):
    S = [0 for _ in range(N+1)]
    for i in range(N):
        S[i+1] = S[i] + A[i]
    
    count_sum = 0
    count = dict()
    for j in range(1, N+1):
        j_mod_k = (S[j] - j) % K  # 0 <= jmodk < K
        if j == 1:
            for i in range(max(0, j-K+1), j):
                # (S[i] - i) % K の値 i_mod_k は、次のループでもほぼ重複する
                # i = max(0, j-K+1) の時の i_mod_k は、i+1で破棄する
                # i = j の時の i_mod_k は、i+1で生成する
                i_mod_k = (S[i] - i) % K
                if i_mod_k in count:
                    count[i_mod_k] += 1
                else:
                    count[i_mod_k] = 1
        else:
            if max(0, j-K+1) != max(0, (j-1)-K+1):
                i = max(0, (j-1)-K+1)
                i_mod_k = (S[i] - i) % K
                if count[i_mod_k] == 1:
                    count.pop(i_mod_k)
                elif count[i_mod_k] > 1:
                    count[i_mod_k] -= 1
            i = j-1
            i_mod_k = (S[i] - i) % K
            if i_mod_k in count:
                count[i_mod_k] += 1
            else:
                count[i_mod_k] = 1
        count_sum += count[j_mod_k] if j_mod_k in count else 0
    return count_sum

if __name__ == '__main__':
    inputs = read()
    print('%d' % solve(*inputs))

Buggy submission (verdict: Runtime Error):
import queue
import bisect

def read():
    N, K = list(map(int, input().strip().split()))
    A = list(map(int, input().strip().split()))
    return N, K, A

def solve(N, K, A):
    S = [0 for _ in range(N+1)]
    T = [0 for _ in range(N+1)]
    # (S[j] - S[i]) % K == j - i
    # (S[j] - j) % K == (S[i] - i) % K and 0 <= (j - i) < K
    # 1 <= j <= N に対して、(S[j] - j) % K == (S[i] - i) % K　を満たす (j - K) < i < j の個数を調べたい
    # j = kのとき i の調べる区間は (k - K) < i < k　つまり C[k] = sum(C[k-K+1], ... , C[k-1])
    # j = k+1のとき i の調べる区間は (k+1 - K) < i < k+1 つまり C[k+1] = sum(C[k-K+2], ... , C[k])
    # j = k+1 では、j = k の結果を用いて C[k+1] = C[k] - C[k-K+1] + C[k]
    for i in range(N):
        S[i+1] = S[i] + A[i]
    
    count_sum = 0
    count = dict()
    for j in range(1, N+1):
        j_mod_k = (S[j] - j) % K  # 0 <= jmodk < K
        if j == 1:
            for i in range(max(0, j-K+1), j):
                # (S[i] - i) % K の値 i_mod_k は、次のループでもほとんど同じ値が使われることに注意
                # i = max(0, j-K+1) の時の i_mod_k は、i+1で破棄する
                # i = j の時の i_mod_k は、i+1で生成する
                i_mod_k = (S[i] - i) % K
                if i_mod_k in count:
                    count[i_mod_k] += 1
                else:
                    count[i_mod_k] = 1
        else:
            # for i in range(j-K+1, j)
            if max(0, j-K+1) != max(0, (j-1)-K+1):
                i = max(0, (j-1)-K+1)
                i_mod_k = (S[i] - i) % K
                if count[i_mod_k] == 1:
                    count.pop(i_mod_k)
                else:
                    count[i_mod_k] -= 1
            i = j-1
            i_mod_k = (S[i] - i) % K
            if i_mod_k in count:
                count[i_mod_k] += 1
            else:
                count[i_mod_k] = 1
        count_sum += count[j_mod_k] if j_mod_k in count else 0
    return count_sum

if __name__ == '__main__':
    inputs = read()
    print('%d' % solve(*inputs))