Problem p01432 (C++) — buggy submission, verdict: Time Limit Exceeded

// OBïÄh2011 Day4 F : eÌ

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#include <vector>

using namespace std;

typedef vector<vector<double>> Matrix;

Matrix mul(Matrix A, Matrix B) {
  int n = A.size();
  Matrix res(n, vector<double>(n, 0));
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      for (int k = 0; k < n; k++)
        res[i][j] += A[i][k] * B[k][j];
  return res;
}

Matrix pow(Matrix A, int p) {
  if (p == 1)
    return A;
  Matrix res = pow(A, p / 2);
  res = mul(res, res);
  if (p % 2)
    res = mul(res, A);
  return res;
}

void gaussianElimination(Matrix &a) {
  int n = a.size();
  for (int i = 0; i < n; i++) {
    int pivot = i;
    for (int j = i + 1; j < n; j++)
      if (abs(a[j][i]) > abs(a[pivot][i]))
        pivot = j;
    swap(a[i], a[pivot]);
    double div = a[i][i];
    for (int j = 0; j <= n; j++)
      a[i][j] /= div;
    for (int j = i + 1; j < n; j++) {
      double mul = a[j][i];
      for (int k = 0; k <= n; k++)
        a[j][k] -= a[i][k] * mul;
    }
  }
  for (int i = n - 1; i >= 0; i--)
    for (int j = 0; j < i; j++)
      a[j][n] -= a[j][i] * a[i][n];
}

int main() {
  int S, N, K;
  double prob[11][100];
  while (cin >> S >> N >> K) {
    S = abs(S);
    if (N == 1) {
      if (S % K == 0)
        cout << S / K << endl;
      else
        cout << -1 << endl;
      continue;
    }
    for (int i = 0; i <= K; i++)
      for (int j = 0; j <= N * K; j++)
        prob[i][j] = 0;
    // KñÅÇê¾¯Ì£iÞ©ÉÂ¢ÄAm¦ðßé
    prob[0][0] = 1.0;
    for (int i = 0; i < K; i++)
      for (int j = 0; j + N <= N * K; j++)
        for (int k = 1; k <= N; k++)
          prob[i + 1][j + k] += prob[i][j] / N;
    // A§ûö®Å[0,N*K-1]ÌÍÍÌúÒlðßé
    vector<vector<double>> a(N * K, vector<double>(N * K + 1, 0));
    a[0][0] = 1;
    for (int i = 1; i < N * K; i++) {
      a[i][i] = a[i][N * K] = 1.0;
      for (int j = 1; j <= N * K; j++) {
        int next = abs(i - j);
        a[i][next] -= prob[K][j];
      }
    }
    gaussianElimination(a);
    if (S < N * K)
      printf("%.6lf\\n", a[S][N * K]);
    // S>=N*KÈçúÒlÍN*KÌQ»®Å\\¹éÌÅAsñÏÅúÒlðßé
    else {
      Matrix A(N * K + 1, vector<double>(N * K + 1, 0));
      for (int i = 0; i < N * K; i++)
        A[0][i] = prob[K][i + 1];
      A[0][N * K] = 1.0;
      for (int i = 1; i < N * K; i++)
        A[i][i - 1] = 1.0;
      A[N * K][N * K] = 1.0;
      A = pow(A, S - N * K + 1);
      double res = A[0][N * K];
      for (int i = 0; i < N * K; i++)
        res += A[0][i] * a[N * K - 1 - i][N * K];
      printf("%.6lf\\n", res);
    }
  }
}

Error: TLE

Fixed submission:

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <iostream>
#include <string.h>
#include <vector>
using namespace std;

const double eps = 1e-8;
typedef vector<double> array;
typedef vector<array> matrix;
int S, N, K;
double prob[11][101];

matrix identity_matrix(int n) {
  matrix A(n, array(n));
  for (int i = 0; i < n; ++i)
    A[i][i] = 1;
  return A;
}

matrix mul(const matrix &A, const matrix &B) {
  matrix C(A.size(), array(B[0].size()));
  for (int i = 0; i < C.size(); ++i)
    for (int j = 0; j < C[i].size(); ++j)
      for (int k = 0; k < A[i].size(); ++k)
        C[i][j] += A[i][k] * B[k][j];
  return C;
}
matrix pow(const matrix &A, int e) {
  return e == 0       ? identity_matrix(A.size())
         : e % 2 == 0 ? pow(mul(A, A), e / 2)
                      : mul(A, pow(A, e - 1));
}

vector<double> gauss_jordan(const matrix &A, const array &b) {
  int n = A.size();
  matrix B(n, array(n + 1));
  for (int i = 0; i < n; i++)
    for (int j = 0; j < n; j++)
      B[i][j] = A[i][j];

  for (int i = 0; i < n; i++)
    B[i][n] = b[i];

  for (int i = 0; i < n; i++) {
    int pivot = i;
    for (int j = i; j < n; j++) {
      if (abs(B[j][i]) > abs(B[pivot][i]))
        pivot = j;
    }
    swap(B[i], B[pivot]);

    if (abs(B[i][i]) < eps) {
      cout << "error\\n";
      return array();
    }

    for (int j = i + 1; j <= n; j++)
      B[i][j] /= B[i][i];
    for (int j = 0; j < n; j++) {
      if (i != j) {
        for (int k = i + 1; k <= n; k++)
          B[j][k] -= B[j][i] * B[i][k];
      }
    }
  }
  array x(n);

  for (int i = 0; i < n; i++)
    x[i] = B[i][n];
  return x;
}

void solve() {

  S = abs(S);
  if (N == 1) {
    if (S % K == 0)
      cout << S / K << endl;
    else
      cout << -1 << endl;
    return;
  }
  memset(prob, 0, sizeof(prob));
  prob[0][0] = 1.0;
  for (int i = 0; i < K; i++)
    for (int j = 0; j + N <= N * K; j++)
      for (int z = 1; z <= N; z++)
        prob[i + 1][j + z] += prob[i][j] / N;

  matrix A(K * N + 1, array(K * N + 1, 0));
  A[0][0] = 1;
  for (int i = 1; i <= K * N; i++) {
    A[i][i] = 1;
    for (int j = 0; j <= K * N; j++) {
      A[i][abs(i - j)] -= prob[K][j];
    }
  }
  array V(K * N + 1, 0);
  V[0] = 0;
  for (int i = 1; i <= K * N; i++)
    V[i] = 1;

  V = gauss_jordan(A, V);

  matrix a(N * K + 1, array(N * K + 1, 0));
  a[N * K][N * K] = a[0][N * K] = 1;
  for (int i = 0; i < N * K; i++)
    a[0][i] = prob[K][i + 1];
  for (int i = 0; i < K * N - 1; i++)
    a[i + 1][i] = 1;

  a = pow(a, S - N * K + 1);
  matrix ans(1, array(N * K + 1));
  double res = a[0][N * K];

  for (int i = 0; i < N * K; i++) {
    res += a[0][i] * V[N * K - i - 1];
  }

  printf("%.6lf\\n", res);
}

int main() {
  while (cin >> S >> N >> K)
    solve();
  return 0;
}