#include <iostream>
#include <vector>
std::vector<std::vector<bool>> visited;
std::vector<int> prefix_sum_a;
std::vector<int> prefix_sum_b;
std::vector<int> prefix_sum_c;
int n;
long long dp(int pos, int len)
{
if (pos > n)
return 0;
if (visited[pos][len])
return 0;
int number_of_a = prefix_sum_a[pos] - prefix_sum_a[pos - len];
int number_of_b = prefix_sum_b[pos] - prefix_sum_b[pos - len];
int number_of_c = prefix_sum_c[pos] - prefix_sum_c[pos - len];
long long result = 0;
if (number_of_a == number_of_b && number_of_b == number_of_c)
{
result = 1ll;
}
else if (number_of_a == 0 && ((number_of_b * number_of_c == 0) || (number_of_b == number_of_c)))
{
result = 1ll;
}
else if (number_of_b == 0 && ((number_of_a * number_of_c == 0) || (number_of_a == number_of_c)))
{
result = 1ll;
}
else if (number_of_c == 0 && ((number_of_b * number_of_a == 0) || (number_of_b == number_of_a)))
{
result = 1ll;
}
visited[pos][len] = true;
/*if (result == 1)
{
printf("%d %d\n", pos - len + 1, pos);
}*/
result += dp(pos + 1, len + 1) + dp(pos + 1, 1);
return result;
}
int main()
{
std::string chuj;
std::cin >> chuj;
n = chuj.size();
prefix_sum_a.resize(n + 1);
prefix_sum_b.resize(n + 1);
prefix_sum_c.resize(n + 1);
/*if (chuj[0] == 'a')
prefix_sum_a[0] = 1;
else if (chuj[0] == 'b')
prefix_sum_b[0] = 1;
else
prefix_sum_c[0] = 1; */
for (int i = 1; i <= n; i++)
{
prefix_sum_a[i] = prefix_sum_a[i - 1];
prefix_sum_b[i] = prefix_sum_b[i - 1];
prefix_sum_c[i] = prefix_sum_c[i - 1];
if (chuj[i - 1] == 'a')
prefix_sum_a[i] += 1;
else if (chuj[i - 1] == 'b')
prefix_sum_b[i] += 1;
else
prefix_sum_c[i] += 1;
}
visited.resize(n + 1);
for (int i = 0; i <= n; i++)
{
visited[i].assign(n + 1, false);
}
long long result = dp(1, 1);
printf("%lld\n", result);
}
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | #include <iostream> #include <vector> std::vector<std::vector<bool>> visited; std::vector<int> prefix_sum_a; std::vector<int> prefix_sum_b; std::vector<int> prefix_sum_c; int n; long long dp(int pos, int len) { if (pos > n) return 0; if (visited[pos][len]) return 0; int number_of_a = prefix_sum_a[pos] - prefix_sum_a[pos - len]; int number_of_b = prefix_sum_b[pos] - prefix_sum_b[pos - len]; int number_of_c = prefix_sum_c[pos] - prefix_sum_c[pos - len]; long long result = 0; if (number_of_a == number_of_b && number_of_b == number_of_c) { result = 1ll; } else if (number_of_a == 0 && ((number_of_b * number_of_c == 0) || (number_of_b == number_of_c))) { result = 1ll; } else if (number_of_b == 0 && ((number_of_a * number_of_c == 0) || (number_of_a == number_of_c))) { result = 1ll; } else if (number_of_c == 0 && ((number_of_b * number_of_a == 0) || (number_of_b == number_of_a))) { result = 1ll; } visited[pos][len] = true; /*if (result == 1) { printf("%d %d\n", pos - len + 1, pos); }*/ result += dp(pos + 1, len + 1) + dp(pos + 1, 1); return result; } int main() { std::string chuj; std::cin >> chuj; n = chuj.size(); prefix_sum_a.resize(n + 1); prefix_sum_b.resize(n + 1); prefix_sum_c.resize(n + 1); /*if (chuj[0] == 'a') prefix_sum_a[0] = 1; else if (chuj[0] == 'b') prefix_sum_b[0] = 1; else prefix_sum_c[0] = 1; */ for (int i = 1; i <= n; i++) { prefix_sum_a[i] = prefix_sum_a[i - 1]; prefix_sum_b[i] = prefix_sum_b[i - 1]; prefix_sum_c[i] = prefix_sum_c[i - 1]; if (chuj[i - 1] == 'a') prefix_sum_a[i] += 1; else if (chuj[i - 1] == 'b') prefix_sum_b[i] += 1; else prefix_sum_c[i] += 1; } visited.resize(n + 1); for (int i = 0; i <= n; i++) { visited[i].assign(n + 1, false); } long long result = dp(1, 1); printf("%lld\n", result); } |
English