\(S=\sum_{i=1}^N A_i\) と定め、\(K'=S-K\) とします。これにより問題は「和が \(K'\) となるような区間の個数を求めよ」となります。

\(A_i\) が正であることから、\(K'<0\) のとき答えは \(0\)、 \(K'=0\) のとき答えは \(1\) です。以下では \(K'>0\) のケースを考えます。

\(A\) の累積和 \(B\) を \(B_0=0\)、 \(B_i=\sum_{j=1}^{i} A_j\) と定めます。このとき、区間和について \(A_l+\dots+A_r=K'\) であることは、\(B_r-B_{l-1}=K'\) であることと同値です。\(l\) を全探索することで問題は「\(B_r=K'+B_{l-1}\) を満たす \(r\) は存在するか？」となります。これは \(B\) を set で持っておくことで \(O(\log N)\) で判定することができるため、全体で \(O(N\log N)\) でこの問題を解くことができます。

尺取法により \(O(N)\) で解く解法もあります。

実装例 (C++)

#include<bits/stdc++.h>
using namespace std;

int main(){
  int n;
  long long k;
  cin >> n >> k;
  vector<int> a(n);
  for(int i=0; i<n; i++) cin >> a[i];
  vector<long long> s(n+1);
  for(int i=0; i<n; i++){
    s[i+1] = s[i] + a[i];
  }

  k = s.back() - k;
  if(k < 0){
    cout << 0 << endl;
    return 0;
  }
  if(k == 0){
    cout << 1 << endl;
    return 0;
  }

  set<long long> ss(s.begin(), s.end());
  int ans = 0;
  for(auto x: s){
    if(ss.find(x+k) != ss.end()){
      ans++;
    }
  }
  cout << ans << endl;
}

なおこの問題のテストケースには、 python の標準の set を hash 衝突により TLE させるものが入っています。 (in93.txt)

hash 衝突を回避するための簡単な手法としては、適当な値と xor を取ってからsetで管理する方法があります。

実装例 (Python)

import random
N, K = map(int, input().split())
A = list(map(int, input().split()))
S = [0]
for a in A:
  S.append(S[-1]+a)

K = S[-1] - K
if K < 0:
  print(0)
  exit()
if K == 0:
  print(1)
  exit()

r = random.randint(1, 10**9)
SS = set(s ^ r for s in S)
ans = 0
for s in S:
  if (s + K) ^ r in SS:
    ans += 1
print(ans)