CSAcademy 23 C. Permutation Matrix ( Periodic, Union Find )

Round #23 (Div. 2 only)

題意：
給一個長度 M 的排列 P，現在要做一個 N * M 的矩陣，第 i 行是 P^i( P )。
問 M 個列分別的數值總和為何。

限制：
1 ≤ N, M ≤ 1e5

解法：
顯然會形成一些循環，用 dfs 找出這些循環並記錄每個節點在每個循環中哪個位置，預處理每個循環的前綴和，就能 O( 1 ) 計算出每個點 ( 對應到列 ) 的數值總和。

時間 / 空間複雜度：
O( N + M )

#include <bits/stdc++.h>
using namespace std;

typedef vector< int > vi;

const int MAXN = ( int ) 1e5;
const int MAXM = ( int ) 1e5;

int N, M;
int P[ MAXM ];

int fa[ MAXM ];

int find( int x ){
  return fa[ x ] == x ? x : fa[ x ] = find( fa[ x ] );
}

void unite( int x, int y ){
  int a = find( x );
  int b = find( y );
  if( a < b ) swap( a, b );
  fa[ a ] = b;
}

int id[ MAXM ];
vi f[ MAXM ];
vi pfx[ MAXM ];

void dfs( int u, int h ){
  if( u == h ) return;
  unite( u, h );
  id[ u ] = f[ h ].size();
  f[ h ].emplace_back( u );
  dfs( P[ u ], h );
}

long long ans[ MAXM ];

signed main(){
  ios::sync_with_stdio( 0 );
  cin >> M >> N;
  for( int i = 0; i < M; ++i ){
    cin >> P[ i ];
    --P[ i ];
    fa[ i ] = i;
  }
  for( int i = 0; i < M; ++i ){
    if( find( i ) != i ) continue;
    f[ i ].emplace_back( i );
    dfs( P[ i ], i );
  }
  for( int i = 0; i < M; ++i ){
    if( find( i ) != i ) continue;
    pfx[ i ] = vi( f[ i ].size() + 1 );
    for( int j = 0; j < f[ i ].size(); ++j ){
      pfx[ i ][ j + 1 ] = pfx[ i ][ j ] + f[ i ][ j ] + 1;
    }
  }
  for( int i = 0; i < M; ++i ){
    int h = find( i );
    int plen = f[ h ].size();
    ans[ i ] += 1LL * N / plen * pfx[ h ][ plen ];
    int n = N % plen;
    int x = id[ P[ i ] ];
    ans[ i ] += pfx[ h ][ min( plen, x + n ) ] - pfx[ h ][ x ];
    ans[ i ] += pfx[ h ][ max( 0, x + n - plen ) ];
  }
  for( int i = 0; i < M; ++i ){
    cout << ans[ i ] << " \n"[ i + 1 == M ];
  }
  return 0;
}