A and B and Interesting Substrings (Codeforces 519D)
Tags: DP, two pointers

A and B are preparing themselves for programming contests.After several years of doing sports programming and solving many problems that require calculating all sorts of abstract objects, A and B also developed rather peculiar tastes.A likes lowercase letters of the Latin alphabet. He has assigned to each letter a number that shows how much he likes that letter (he has assigned negative numbers to the letters he dislikes). B likes substrings. He especially likes the ones that start and end with the same letter (their length must exceed one).Also, A and B have a string s. Now they are trying to find out how many substrings t of a string s are interesting to B (that is, t starts and ends with the same letter and its length is larger than one), and also the sum of values of all letters (assigned by A), except for the first and the last one is equal to zero.Naturally, A and B have quickly found the number of substrings t that are interesting to them. Can you do it? 

Input: The first line contains 26 integers xa, xb, ..., xz ( - 105 ≤ xi ≤ 105) — the value assigned to letters a, b, c, ..., z respectively.The second line contains string s of length between 1 and 105 characters, consisting of Lating lowercase letters— the string for which you need to calculate the answer. 

Output: Print the answer to the problem. 

Note: In the first sample test strings satisfying the condition above are abca and bcab.In the second sample test strings satisfying the condition above are two occurences of aa.

Sample Inputs:
1 1 -1 1 1 1 1 1 1 1 1 1 1 1 1 7 1 1 1 8 1 1 1 1 1 1
xabcab

1 1 -1 1 1 1 1 1 1 1 1 1 1 1 1 7 1 1 1 8 1 1 1 1 1 1
aaa

Sample Outputs:
2

2

Solution (Python):
def solve(x, s):
    n = len(s)
    cnt = [0] * 26
    ans = 0
    for i in range(26):
        cnt[i] = x[i]
    sum_arr = [0] * (n + 1)
    pos = [[] for _ in range(26)]
    for i in range(1, n + 1):
        sum_arr[i] = sum_arr[i - 1] + cnt[ord(s[i - 1]) - ord('a')]
        pos[ord(s[i - 1]) - ord('a')].append(i)
    for cc in range(26):
        mp = {}
        for p in pos[cc]:
            ans += mp.get(sum_arr[p - 1], 0)
            mp[sum_arr[p]] = mp.get(sum_arr[p], 0) + 1
    return ans
x = list(map(int, input().strip().split()))
s = input().strip()
res = solve(x, s)
print(res)

Solution (Python):
import sys
input = sys.stdin.readline

def inp():
    return int(input())

def inlt():
    return list(map(int, input().split()))

def insr():
    s = input()
    return list(s[:len(s) - 1])

def invr():
    return map(int, input().split())

def A_and_B_and_Interesting_Substrings2():
    value_sequence = inlt()
    s = insr()
    n = len(s)
    char_pos_dict = {}
    prefix_sum = []
    running_sum = 0
    for (index, char) in enumerate(s):
        ascii_code = ord(char) - ord('a')
        running_sum += value_sequence[ascii_code]
        prefix_sum.append(running_sum)
        if char not in char_pos_dict.keys():
            char_pos_dict[char] = []
        char_pos_dict[char].append(index)
    max_prefix_sum = max(prefix_sum)
    total_interesting_substrings = 0
    for ascii_code in range(ord('a'), ord('a') + 26):
        char = chr(ascii_code)
        if char in char_pos_dict.keys() and len(char_pos_dict[char]) > 1:
            pos_list = char_pos_dict[char]
            count_of_prefix_sums_dict = {}
            for (index, pos) in enumerate(pos_list):
                if index == 0:
                    p_sum = prefix_sum[pos]
                    if p_sum not in count_of_prefix_sums_dict.keys():
                        count_of_prefix_sums_dict[p_sum] = 0
                    count_of_prefix_sums_dict[p_sum] += 1
                else:
                    p_sum = prefix_sum[pos]
                    p_sum_target = prefix_sum[pos - 1]
                    if p_sum_target in count_of_prefix_sums_dict.keys():
                        total_interesting_substrings += count_of_prefix_sums_dict[p_sum_target]
                    if p_sum not in count_of_prefix_sums_dict.keys():
                        count_of_prefix_sums_dict[p_sum] = 0
                    count_of_prefix_sums_dict[p_sum] += 1
    print(total_interesting_substrings)
    return
A_and_B_and_Interesting_Substrings2()

Solution (Python):
if __name__ == '__main__':
    X = [int(i) for i in input().split()]
    s = input()
    arr = []
    unicode_a = ord('a')
    for i in range(26):
        arr.append({})
    sm = 0
    k = 0
    for char in s:
        n = ord(char) - unicode_a
        k += arr[n].get(sm, 0)
        sm += X[n]
        arr[n][sm] = arr[n].get(sm, 0) + 1
    print(k)