#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); } |