G - Takahashi And Pass-The-Ball Game Editorial by en_translator
First, consider the state where the operation is actually performed once, to assume that between operation is performed between \(0\) and \(K-1\) more times. We seek for the answer multiplied by \(K\) and finally divide it by \(K\). Then, this problem asks to perform the operation \(K\) times and find for each Takahashi the sum of the number of balls after each operation.
By adding the number of balls, two consecutive operations can be considered at once, so this problem can be solved by the doubling technique, solving it for the halved \(K\). If \(K\) is odd, we can consider the initial state and the state after the operation is performed once, and do the same process.
Mathematical explanation
Define the following operations for length-\(N\) sequences \(S=(S _ 1,S _ 2,\ldots,S _ N)\) and \(T=(T _ 1,T _ 2,\ldots,T _ N)\) consisting of \(\{1,2,\ldots,N\}\), length-\(N\) sequences \(\bm{x}=(x _ 1,x _ 2,\ldots,x _ N)\) and \(\bm{y}=(y _ 1,y _ 2,\ldots,y _ N)\) of \(\mathbb{F}_p\), and a non-negative integer \(n\):
\[\bm x+\bm y=(x _ 1+y _ 1,x _ 2+y _ 2,\ldots,x _ N+y _ N)\] \[S\times\bm x=\left(\sum _ {S _ j=i}x _ j\right) _ {i=1,2,\ldots,N}\] \[S\times T=(S _ {T _ 1},S _ {T _ 2},\ldots,S _ {T _ N})\] \[S ^ n=\left\{\begin{matrix}(1,2,\ldots,N)&(n=0)\\S ^ {n-1}\times S&(n\gt 0)\end{matrix}\right.\]
The first operation find the sum of two states of balls. The second operation represents “the state of balls right after passing balls as \(i\to S _ i\) when the \(i\)-th Takahashi has \(x _ i\) balls.” The third operation represents “after passing balls as \(i\to T _ i\) and then \(i\to S _ i\), the \(i\)-th Takahashi’s balls go to which Takahashi?” The fourth operation represents “after passing balls as \(i\to S _ i\) \(n\) times, the \(i\)-th Takahashi’s balls go to which Takahashi?”
The first three operations can be computed in a total of \(O(N)\) time.
Using these operations, the answer can be expressed as follows:
\[\bm b+(A ^ 1\times\bm b)+\cdots+(A ^ {K-1}\times\bm b)\]
Here, \(A\) is \((A _ i)\) given in the input, and \(\bm b=(b _ i) _ i\) it the number of balls that Takahashi has after one operation.
Consider how to evaluate the function \(f(S,\bm x,k)\coloneqq\bm x+(S\times\bm x)+\cdots+(S ^ {k-1}\times\bm x)\). What we want is \(f(A,\bm b,K)\).
\[\begin{aligned} &f(S,\bm x,k)\\ ={}&\bm x+(S\times\bm x)+\cdots+(S ^ {k-1}\times\bm x)\\ ={}&\bm x+(S\times\bm x)+(S\times(S\times\bm x))+\cdots+(S^{k-2}\times(S\times\bm x))\\ ={}&\bm x+f(S,S\times\bm x,k-1) \end{aligned}\]
\[\begin{aligned} &f(S,\bm x,2k)\\ ={}&\bm x+(S\times\bm x)+\cdots+(S ^ {2k-1}\times\bm x)\\ ={}&(\bm x+S\times\bm x)+(S ^ 2\times(\bm x+S\times\bm x))+\cdots+((S ^ 2) ^ {k - 1}\times(\bm x+S\times\bm x))\\ ={}&f(S ^ 2,\bm x+S\times\bm x,k) \end{aligned}\]
Here, we use the properties of operations like \(S\times(\bm x+\bm y)=(S\times\bm x)+(S\times\bm y)\), \((A\times B)\times\bm x=A\times(B\times\bm x)\), and so on.
Therefore, it can be boiled down to a problem with one size smaller and half the size, so the problem can be solved in a total of \(O(N\log K)\) time.
By implementing this, the problem can be solved in a total of \(O(N\log K+\log\operatorname{mod})\) time.
The formulation is in a form of tail-recursion, so it can also be written with a while statement.
Also, one can use the property of a functional graph (a graph with \(N\) directed edges \(i\to A _ i\)) to solve the problem in a total of \(O(N+\log\operatorname{mod})\) time.
#include <bits/extc++.h>
#include <atcoder/modint>
template<int m>
std::ostream& operator<<(std::ostream& os, const atcoder::static_modint<m>& mint){
os << mint.val();
return os;
}
int main() {
using namespace std;
using modint = atcoder::static_modint<998244353>;
unsigned long N, K;
cin >> N >> K;
const auto K_inv{modint{K}.inv()};
vector<unsigned long> A(N);
for(auto&& a : A){
cin >> a;
--a;
}
vector<modint> ball(N);
for(auto&& b : ball){
int x;
cin >> x;
b = modint::raw(x);
}
// Composition of operations i → B[i] → A[B[i]]
const auto compose{[N](auto&& A, const auto& B){
vector<unsigned long> result(N);
for(unsigned long i{}; i < N; ++i)result[i] = A[B[i]];
A = result;
}};
// Move the balls
const auto apply{[N](const auto& A, const auto& x){
vector<modint> result(N);
for(unsigned long i{}; i < N; ++i)result[A[i]] += x[i];
return result;
}};
const auto add{[N](auto&& x, const auto& y){
for(unsigned long i{}; i < N; ++i)x[i] += y[i];
}};
ball = apply(A, ball); // Perform the operation once
vector<modint> ans(N);
while(K){
if(K & 1){ // If K is odd, consider the first and the remaining
add(ans, ball);
ball = apply(A, ball);
}
// Put the two together
add(ball, apply(A, ball));
compose(A, A);
K /= 2;
}
// Multiply by 1/K to find the answer
for(const auto& x : ans)cout << x * K_inv << " ";
cout << endl;
return 0;
}
posted:
last update: