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)，有

[egin{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} ]

然而

[egin{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;
}