PROBLEM LINK:
Practice
Div-3 Contest
Div-2 Contest
Div-1 Contest
Author: Vsevolod Lepeshov
Tester: Felipe Mota and Radoslav Dimitrov
Editorialist: Vsevolod Lepeshov
DIFFICULTY:
Easy
PREREQUISITES:
Strings, Bruteforce
PROBLEM:
There are N funny words s_1, s_2, \dots, s_N.
You have to calculate the number of ways to choose two words such that they are not funny, but if we swap the first letters in them, the resulting words are funny.
QUICK EXPLANATION:
Divide all words in two parts - first letter and whole word except first letter.
Numerate first letters with numbers from 0 to 25, and whole word except first letter from 0 to n using coordinate compression or map or hashes, which works in O(|s_i| * n * log(n)).
Iterate over first letter in both words (26 * 26 options), let’s call f_1 - first letter in first word, f_2 - first letter in second word.
Iterate over n “compressed” suffixes. Let’s calculate cnt_{01} - number of suffixes such that there is not funny word which starts with f_1, but there is funny word that starts with f_2, and cnt_{10} - number of suffixes such that there is not funny word which starts with f_2, but there is funny word that starts with f_1.
The number of good team names with this first letters is cnt_{01} * cnt_{10}, so the answer is the sum of this values over all first letters in both words.
Total asymptotics: O(|s_i| * n * log(n) + 26 * 26 * n)
EXPLANATION:
Each word contains two parts - first letter and “whole word except first letter”. Let’s call “whole word except first letter” - suffix.
For example word “abba” have first letter “a” and suffix “bba”.
Let’s call f_i - number of first letter of s_i. 0 \le f_i \lt 26.
Note that we do not need the suffixs of the words themselves to solve the problem. The only thing we need is to quickly understand whether the suffixes of some words are equal.
So, we will construct array c_1, \dots, c_n such that the suffixes of s_i and s_j are equal if and only if c_i and c_j are equal.
One way to build such array is to sort all the words, and for the word s_i, set c_i as “The number of words lexicographically smaller than s_i”. This technique is called coordinate compression.
The second way: just use the map<string, int> (or dict in Python) which return c_i by suffix of s_i. Iterate over all suffixes: if suffix of s_i already in map you know c_i for it, else you can assign new c_i for this suffix and add it to map.
After that each word can be coded by pair - (f_i, c_i).
Let’s construct boolean 2d array isword[26][n]. isword[i][j] means - “Is there a word that coded by pair (i, j).” So, we will set isword[f_i][s_i] = true for all i from 1 to n, and false for other cells.
We want to understand when (X_1, Y_1), (X_2, Y_2) is good teamname. If it is good teamname - words (X_1, Y_1) and (X_2, Y_2) are not funny, but words (X_2, Y_1) and (X_1, Y_2) are funny. In other words:
isword[X_1][Y_1] = false
isword[X_2][Y_2] = false
isword[X_1][Y_2] = true
isword[X_2][Y_1] = true
Let’s iterate over first letters - X_1 and X_2. We want to count number of good teamnames such that first words starts with X_1 and second with X_2. We will sum this values over all X_1 and X_2 to get final answer.
Iterate over all suffixes and count number of Y_1 such that
isword[X_1][Y_1] = false
isword[X_2][Y_1] = true
Let’s call this number cnt_{01}
Also let’s count number of Y_2 such that
isword[X_2][Y_2] = false
isword[X_1][Y_2] = true
Let’s call this number cnt_{10}
So (X_1, Y_1), (X_2, Y_2) is good teamname if and only if Y_1 is in one of cnt_{01} values, and Y_2 in one of cnt_{10} values. So the total number is cnt_{01} * cnt_{10}. Sum this values over all X_1, X_2 to get final answer.
ALTERNATE EXPLANATION:
SOLUTIONS:
Setter's Solution
#include <bits/stdc++.h>
#define pb push_back
#define ld long double
using namespace std;
signed main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
cout.tie(nullptr);
int t;
cin >> t;
while (t--) {
int n;
cin >> n;
vector<string> S(n);
for (int i = 0; i < n; i++) {
cin >> S[i];
}
vector<pair<string, int>> suffixes;
// Actually, we can divide string in 2 parts - "first letter" and "all string except first letter"
for (int i = 0; i < n; i++) {
string suf = "";
for (int x = 1; x < (int) S[i].size(); x++) {
suf += S[i][x];
}
// suf is "all string except first letter"
suffixes.pb({suf, i});
}
// Now we will compress all suffixes
// We want to make array which have property:
// "suffixes i and j are equal if and only if compress_i == compress_j"
// It is simple to construct this array after sorting all suffixes
sort(suffixes.begin(), suffixes.end());
vector<int> compress(n);
for (int i = 1; i < n; i++) {
if (suffixes[i].first == suffixes[i - 1].first) {
compress[suffixes[i].second] = compress[suffixes[i - 1].second];
} else {
compress[suffixes[i].second] = 1 + compress[suffixes[i - 1].second];
}
}
// is_word[i][j] means "is there a funny word with first letter number i, and suffix with "compress number" j"
vector<vector<bool>> is_word(26, vector<bool>(n, false));
for (int i = 0; i < n; i++) {
// for i-th word we set needed is_word cell to true
is_word[S[i][0] - 'a'][compress[i]] = true;
}
// Now we will bruteforse first letters in both words in good team name
int ans = 0;
for (int i = 0; i < 26; i++) {
for (int j = 0; j < 26; j++) {
// We will iterate over all n possible suffixes
int cnt_10 = 0, cnt_01 = 0;
for (int x = 0; x < n; x++) {
if (is_word[i][x] && !is_word[j][x]) {
cnt_10++;
} else if (!is_word[i][x] && is_word[j][x]) {
cnt_01++;
}
}
// We can combine each "01" suffix with each "10" suffix, so in total cnt_10 * cnt_01
ans += cnt_10 * cnt_01;
}
}
cout << ans << '\n';
}
}
