We can notice that our problem is equivalent to the following one:
Let’s draw square field of size N x N, then we need to mark N cells, such that all of them are placed in different rows and columns and none of them is placed in bottom-left staircase of size A x A and upper-right staircase of size N - A x N - A.
Let’s start from slow solution:
Suppose we will choose cells, going from up to down. For this we will write dp[i][k] ~— number of ways if we’ve considered first i rows and we have chosen k cells from the upper-left rectangle(of size N - A x A), assuming that we don’t care which cells exactly we will take from upper-left rectangle.
So, we have two types of transition:
- We take cell from the left part:
dp[i + 1][k + 1] += dp[i][k]
- We take cell from the right part. Since in this case, cells from i - k columns are taken, we have k options:
dp[i + 1][k] += k \cdot dp[i][k]
So we will compute this table for the first N - A rows and last A rows(in this case, going from down to up and maintaing number of cells taken from the bottom-right rectangle).
But then, if we will fix k ~— how many cells we have taken from the upper-left rectangle, we can notice that we should take exactly k cells from the bottom-right rectangle too. Also, cells, which we’ve taken not from rectangles, define k columns for the upper-left and bottom-right rectangles. Since while computing our dp, we didn’t care which columns will be used in our rectangles, we should multiply each way by (k!)^2. Finally, answer for the fixed k is equal to:
dp[A][k] \cdot dp[N - A][k] \cdot (k!)^2.
Summing over all k will give us the answer.
To optimize it, let’s speed up way of calculating dp[i][\cdot]. Actually, we can notice that described transitions are the same as transitions for computing stirling number of the second type(or we can just print few small values and search it in oeis).
So, for fixed value of n we need to compute S(n, k) for 1 \le k \le A.
It’s quite well-known, we can do it using the explicit formula:
S(n, k) = \frac{1}{k!}\sum\limits_{j=0}^{k}(-1)^{k-j}\binom{k}{j}j^n
We can optimize it, using fft for polynomials with coefficients:
a_i = \frac{i^n}{i!}, b_i = (-1)^i \frac{1}{i!}.
Overall, the solution works in O(Alog(A))