[Algorithm] BOJ 28212. Classical Summation Problem

https://www.acmicpc.net/problem/28212

 

 \(1\) 이상 \(n\) 이하의 자연수 \(k\)개로 이루어진 임의의 쌍 \((a_1, ..., a_k)\)에 대하여 함수 \(y=|x-a_1|+...+|x-a_k|\)가 최소가 되는 \(x\)의 값들 중 최솟값을 모든 가능한 \((a_1, ..., a_k)\) 쌍에 대하여 구한 후 그 합을 \(998244353\)으로 나눈 나머지를 출력하는 문제입니다. 시간 제한은 \(2\)초, 메모리 제한은 \(1024\)MB입니다. 

 

 함수 \(y=|x-a_1|+...+|x-a_k|\)가 최소가 되는 \(x\)의 값들 중 최솟값이 \(i\)라는 것은, \(a_1, ..., a_k\)을 오름차순으로 정렬했을 때 \(1\) 이상 \(i-1\) 이하의 자연수의 개수가 \(\lfloor \frac{k+1}{2} \rfloor\)개 미만이고, \(1\) 이상 \(i\) 이하의 자연수의 개수가 최초로 \(\lfloor \frac{k+1}{2} \rfloor\)개 이상이 되었다는 것입니다.

 

 다르게 말하면, 함수 \(y=|x-a_1|+...+|x-a_k|\)가 최소가 되는 \(x\)의 값들 중 최솟값이 \(i\) 이상이라는 것은, \(a_1, ..., a_k\)을 오름차순으로 정렬했을 때 \(1\) 이상 \(i-1\) 이하의 자연수의 개수가 \(\lfloor \frac{k+1}{2} \rfloor\)개 미만이라는 것입니다.

 

 함수 \(y=|x-a_1|+...+|x-a_k|\)가 최소가 되는 \(x\)의 값들 중 최솟값이 \(i\)인 \((a_1, ..., a_k)\) 쌍의 개수를 \(N(X=i)\), 최솟값이 \(i\) 이상인 \((a_1, ..., a_k)\) 쌍의 개수를 \(N(X\geq i)\)로 표기합시다. 그러면 문제의 답은 다음과 같습니다. \[\sum_{i=1}^{n}iN(X=i)=\sum_{i=1}^{n}N(X\geq i) \mod 998244353\]

 

 \(N(X\geq i)\)는 \(1, ..., i-1\) 중에서 중복을 허용하여 \(j\)개를 뽑고 (단, \(0 \leq j \leq \lfloor \frac{k+1}{2} \rfloor-1\)), \(i, ..., n\) 중에서 중복을 허용하여 \(k-j\)개를 뽑은 후, 이들로 \((a_1, ..., a_k)\)를 구성하는 경우의 수를 모두 더하여 구할 수 있습니다. 단, 이때 \(0^0=1\)으로 생각합니다. \[N(X \geq i)=\sum_{j=0}^{\lfloor \frac{k-1}{2} \rfloor}\binom{k}{j}(i-1)^j(n-i+1)^{k-j}\]

 

 이를 종합하면 답은 다음과 같습니다. \begin{align}& \sum_{i=1}^{n}\sum_{j=0}^{\lfloor \frac{k-1}{2} \rfloor}\binom{k}{j}(i-1)^j(n-i+1)^{k-j} \\ &= n^k + \sum_{i=1}^{n-1}\sum_{j=0}^{\lfloor \frac{k-1}{2} \rfloor}\binom{k}{j}i^j(n-i)^{k-j} \\ &= n^k + \frac{1}{2}\sum_{i=1}^{n-1}\sum_{j=0}^{\lfloor \frac{k-1}{2} \rfloor}[\binom{k}{j}i^j(n-i)^{k-j} + \binom{k}{k-j}i^{k-j}(n-i)^j] \\ &= \begin{cases} n^k + \frac{1}{2}\sum_{i=1}^{n-1}\sum_{j=0}^{k}\binom{k}{j}i^j(n-i)^{k-j} & \text{ if } k \text{ is odd} \\  n^k + \frac{1}{2}\sum_{i=1}^{n-1}[(\sum_{j=0}^{k}\binom{k}{j}i^j(n-i)^{k-j})-\binom{k}{k/2}i^{k/2}(n-i)^{k/2}] & \text{ if } k \text{ is even} \end{cases} \\ &= \begin{cases} n^k + \frac{1}{2}(n-1)n^k & \text{ if } k \text{ is odd} \\  n^k + \frac{1}{2}\sum_{i=1}^{n-1}[n^k-\binom{k}{k/2}i^{k/2}(n-i)^{k/2}] & \text{ if } k \text{ is even} \end{cases} \\ &= \frac{1}{2}(n+1)n^k - \begin{cases}  0 & \text{ if } k \text{ is odd} \\  \frac{1}{2}\binom{k}{k/2}\sum_{i=1}^{n-1}i^{k/2}(n-i)^{k/2} & \text{ if } k \text{ is even} \end{cases} \mod 998244353 \end{align}

 

 시간복잡도는 \(O(n\lg k)\) 입니다.

 

 코드는 다음과 같습니다.

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define mod 998244353

ll f[1000010], invf[1000010];

ll power(ll a, ll b)
{
    ll ret=1;
    while(b)
    {
        if(b&1) ret=ret*a%mod;
        a=a*a%mod;
        b>>=1;
    }
    return ret;
}

ll inv(ll a)
{
    return power(a, mod-2);
}

ll binom(ll n, ll k)
{
    return f[n]*invf[k]%mod*invf[n-k]%mod;
}

int main()
{
    f[0]=1;
    for(ll i=1;i<=1000000;i++) f[i]=f[i-1]*i%mod;
    invf[1000000]=power(f[1000000], mod-2);
    for(ll i=1000000-1;i>=0;i--) invf[i]=invf[i+1]*(i+1)%mod;
    assert(invf[0]==1);
    
    ll n, k;
    cin>>n>>k;
    ll ans=(n+1)*power(n, k)%mod*inv(2)%mod;
    if(k&1) cout<<ans;
    else
    {
        ll tmp=0;
        for(ll i=1;i<=n-1;i++)
        {
            tmp=(tmp+power(i*(n-i)%mod, k/2))%mod;
        }
        tmp=tmp*binom(k, k/2)%mod*inv(2)%mod;
        cout<<(ans-tmp+mod)%mod;
    }
}