519D - A and B and Interesting Substrings

Solution:

Another cool problem. This problem can be solved using a dynamic programming approach.

Firstly let sum[i] be the cumulative sum of the values assigned to the letters in s[0 .. i]. This will allow us to find the sum of values between two letters in O(1).

There are 26 letters to consider, but let's focus on only one of them first. For a given letter c, we firstly want to know in which positions letter c occurs in the string. Between two letters c with index i and j, what is the condition such that the sum of values between the two letters be 0? Answer is when sum[i] == sum[j-1], and simply that! So what we can do is to maintain a variable cnt that counts the total of pairs that satisfy that condition, and a map M which keeps track the number of occurrence of each possible sum. Then for i = 0 to last letter c :

1. we check if sum[i-1] is in M, and if so we add the current value of M[sum[i-1]] to cnt.

2. we increment the counter of sum[i] in M by one.

Finally we do the same to all letters and total up the result to arrive at the answer.

#include <iostream> #include <cstdio> #include <algorithm> #include <vector> #include <string> #include <map> using namespace std; vector<int> t[30]; string s; long long sum[100005]; int val[30]; int N; int main(){ for(int i=0;i<26;++i){ scanf("%d",&val[i]); } cin >> s; N = s.size(); sum[0] = 0; for(int i=0;i<N;++i){ sum[i+1] = sum[i] + val[s[i]-'a']; t[s[i]-'a'].push_back(i+1); } long long ans = 0; for(int k=0;k<26;++k){ long long tot = 0; map<long long,int> tmp; for(int i=0;i<t[k].size();++i){ if(tmp.count(sum[t[k][i]-1])) { tot += tmp[sum[t[k][i]-1]]; } if(tmp.count(sum[t[k][i]])) { tmp[sum[t[k][i]]]++; } else { tmp[sum[t[k][i]]] = 1; } } ans += tot; } cout << ans << endl; return 0; }