Johnny and Grandmaster

[CodeForces - 1361B] Johnny and Grandmaster

将\(n\)个数：\(p^{k_1}, p^{k_2},...,p^{k_n}\)，分成互不相交的两组，这两组内的元素和分别记为\(sum_1\)和\(sum_2\)，求\(|sum_1 - sum_2|_{min} \% mod\)是多少？

\(1 \leq p, n, k_i \leq 1e6\)，\(mod = 1e9 + 7\)

题解

容易想到的是，\(sum_1\)和\(sum_2\)分别就是两个以\(p\)进制表示的整数(只是某些位置上可能重复)。先将\(n\)个数按降序排列，为了让它们的差尽量小，将最大的\(p^{k_0}\)放入\(sum_1\)，\(p^{k_1}\)放入\(sum_2\)，如果\(sum_1 > sum_2\)，则把\(p^{k_2}\)放入\(sum_2\)，如果依然小于，重复这个过程，直到\(sum_1 \leq sum_2\)

定理 \(\exists\) \(i \in N\)使得\(p^{k_0} = p^{k_1} +p^{k_2} +p^{k_3} +\cdots + p^{k_i}\)

证明：对于任意一个长度为\(n\)的上述数列，如果\(p^{k_0} > p^{k_1} + p^{k_2} + \cdots + p^{k_{n-1}}\)，那么再增加 \(t\) 个 \(p^1\)，就可以使得等号成立（进制的性质），既\(i = n - 1 + t\)；如果\(p^{k_0} = p^{k_1} + p^{k_2} + \cdots + p^{k_{n-1}}\)，显然\(i = n - 1\)；如果\(p^{k_0} < p^{k_1} + p^{k_2} + \cdots + p^{k_{n-1}}\)，则一定会有一个\(i <n - 1\)，使得等号成立，理由如下，假设不存在这样的\(i\)，既存在某个\(j < n - 1\)，有

\[\begin{align} p^{k_0} &> p^{k_1} + p^{k_2} + \cdots + p^{k_j} &(1)\\ p^{k_0} &< p^{k_1} + p^{k_2} + \cdots + p^{k_j} + p^{k_{j+1}} &(2) \end{align} \]

然而

\[\begin{align} p^{k_0} - p^{k_1} &= (p^{k_0 -k_1}-1)p^{k_1} \\ p^{k_0} - p^{k_1} - p^{k_2} &=((p^{k_0 -k_1}-1)p^{k_1 - k_2}-1)p^{k_2} \\ &\cdots \\ p^{k_0} - p^{k_1} - p^{k_2} - \cdots - p^{k_j} &= coe*p^{k_j} \\ \end{align} \]

\(coe\)代表系数；上述推导过程说明式(1)的右边再增加 \(coe(coe \geq 1)\) 个\(p^{k_j}\)，就能使得\(p^{k_0} = p^{k_1} + p^{k_2} + \cdots + p^{k_j} + coe*p^{k_j}\)，又因为按降序排列，故\(p^{k_{j+1}} \leq p^{k_j}\)，所以

\[p^{k_1} + p^{k_2} + \cdots + p^{k_j} + p^{k_{j+1}} \leq p^{k_1} + p^{k_2} + \cdots + p^{k_j} + coe*p^{k_j} =p^{k_0} \]

显然这与式(2)矛盾，故假设不成立，既 \(\exists\) \(i < n - 1\)，使得等号成立。

综上所述，定理成立。

那么怎么找到这些 \(i\) 呢？举个例子：\(p = 2\)，\(k = \{3,2,1\}\)，\(mod = 7\)

质朴的做法是先用快速幂求出来每个数，再来贪心：\(sum_1 = 2^3 \% mod = 1\)，\(sum_2 = 2^2 \% mod = 4\)，显然\(sum_1 < sum_2\)，故\(sum_1 = 1 + 2^1 \% mod = 3\)，\(|sum_1-sum_2| \% mod = |3 - 4| \% mod = 1\)，稍微用笔算一下就会发现，存在更优解\(|sum_1 - sum_2|_{min} \% mod = 2\)。因为在实际例子中，\(p^{k_i}\)可能会爆\(long\) \(long\)，快速幂必须要取模才能计算。而上述栗子说明：取模操作导致无法判断大小关系。

实际上，观察定理的推导过程我们会发现，可以递推计算系数\(coe\)，并且如果\(coe >1e6\)，说明即使把剩下的数都加入同一组，也不能使等号关系成立（既不用快速幂计算系数从而避免了取模的问题），而这部分才具有真正的贡献。当\(coe = 0\)时，说明\(current\_sum_1 = current\_sum_2\)。

代码

#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;

int n, p;
const int mod = 1e9 + 7;

long long fast_pow(int a, int x) {
    long long ans = 1;
    while(x) {
        if (x & 1) ans = ans * a % mod;
        x >>= 1;
        a = 1ll * a * a % mod;
    }
    return ans;
}


int main()
{
    int cas;
    cin >> cas;

    vector<int> num;
    while(cas--) {
        cin >> n >> p;

        num.clear();
        for (int i = 0; i < n; i++) {
            int x;
            cin >> x;
            num.push_back(x);
        }

        if (p == 1) {
            cout << (n & 1) << endl;
            continue;
        }

        sort(num.begin(), num.end());

        long long cur_coe = 0;
        int pos = n - 1;                 // 具有真正贡献的那部分的第一个元素 

        for (int i = n - 1; i >= 0; i--) {
            if (cur_coe == 0) {         // 说明前面已经分成的两组元素和相等，
                pos = i;
                cur_coe = 1;
                continue;
            }
            int delta = num[i + 1] - num[i];

            long long tp = 1;
            for (int j = 0; j < delta; j++) {
                tp = tp * p;
                if (tp > 1000000) goto label;
            }

            cur_coe = cur_coe * tp - 1;     // 递推计算系数 coe
            if (cur_coe > 1000000) break;
        }

        label:;

        long long ans = fast_pow(p, num[pos]);
        for (int i = 0; i < pos; i++) {
            ans = (ans - fast_pow(p, num[i]) + mod) % mod;
        }

        cout << (ans + mod) % mod << endl;
    }

    return 0;
}