ABC332 E - Lucky bag 半分全列挙解法

AtCoder Beginner Contest 332 E - Lucky bagを半分全列挙で解いたので解法を書いておきます。

問題

解法

半分に分けて割り当て方を列挙する

すべてのグッズを振り分ける方法を考えると、 $N \le 15$ とはいえ間に合いそうにないです。
（実際に1382958545通りあります。）
そこで、グッズを半分に分けてから振り分け方を考えます。
ここはDFSやBFSを用いればできます。

実装例

C++

vector<vector<ll>> a, b;
vector<ll> v;
auto dfs = [&](auto self, vector<vector<ll>> &vec, ll start, ll end,
               ll index) -> void {
    if (index == end) {
        vec.push_back(v);
        return;
    }
    REP(i, v.size()) {
        v[i] += w[index];
        self(self, vec, start, end, index + 1);
        v[i] -= w[index];
    }
    if (ll(v.size()) < d) {
        print(w[index]);
        v.push_back(w[index]);
        self(self, vec, start, end, index + 1);
        v.pop_back();
    }
};
dfs(dfs, a, 0, n / 2, 0);
dfs(dfs, b, n / 2, n, n / 2);

Python

def divide_team(start, end):
    q = [(start, [])]
    while q:
        i, a = q.pop()
        if i == end:
            yield a
            continue
        for j in range(len(a)):
            b = a[:]
            b[j] += w[i]
            q.append((i + 1, b))
        if len(a) < d:
            q.append((i + 1, a + [w[i]]))

a = tuple(divide_team(0, n // 2))
b = tuple(divide_team(n // 2, n))

マージする

こうして考えた振り分け方ですが、グッズ7個の分け方が最大877通り、8個の分け方が最大4140通りであるためその組み合わせは $877 \times 4140 = 3630780$ 通りとなり、うまくマージすれば制限時間に間に合います。
そのマージの方法ですが $D$ 個の福袋が左右一列に並んでいるとして、前半は $W$ の合計値の大きいものを左から割り当てる、後半は大きいものを右から割り当てるとすれば最善です。

例） $D = 3$ で前半の割り当て方が $\{3, 5\}$ 、後半が $\{3, 6, 3\}$ のとき
前半は $\{5, 3, 0\}$ 、後半は $\{3, 3, 6\}$ と割り当て、福袋の重さが $\{8, 6, 6\}$ となるときが最善となる。

証明(クリックで表示)

降順の配列 $A = (A_{1}, A_{2}, ..., A_{D})$ と昇順の配列 $B = (B_{1}, B_{2}, ..., B_{D})$ があり、 $C = (A_{1} + B_{1}, A_{2} + B_{2}, ..., A_{D} + B_{D})$ とする。
ここで $B_i, B_j ( i \lt j )$ を入れ替えたときの $C$ の分散の変動を考える。
$A_j = A_i + A_{diff}, B_i = B_j + B_{diff} ( A_{diff}, B_{diff} \ge 0 )$ 、 $C$ の平均を $\bar{c}$ とすると、分散の変動値は
$\{ (A_i + B_j - \bar{c}) ^ 2 + (A_j + B_i - \bar{c}) ^ 2 \} - \{ (A_i + B_i - \bar{c}) ^ 2 + (A_j + B_j - \bar{c}) ^ 2 \}$
$= \{ (A_i + B_j - \bar{c}) ^ 2 + (A_i + A_{diff} + B_j + B_{diff} - \bar{c}) ^ 2 \} - \{ (A_i + B_j + B_{diff} - \bar{c}) ^ 2 + (A_i + A_{diff} + B_j - \bar{c}) ^ 2 \}$
$= 2 A_{diff} B_{diff} \ge 0$
よって分散が減少することはない。
$A$ の要素を入れ替えたときも同様に分散が減少することはなく、 $A$ が降順、 $B$ が昇順のときに分散の最小値を取る。(証明終)

まず割り当て方の候補はあらかじめソートしておき、あとは前後から割り当てていけば解けます。

実装例

C++

for (auto &&i : a) {
    sort(i.rbegin(), i.rend());
}
for (auto &&i : b) {
    sort(i.rbegin(), i.rend());
}
double ans = 1e18;
double ave_w = accumulate(w.begin(), w.end(), 0.0) / d;
for (const auto &i : a)
    for (const auto &j : b) {
        vector<ll> bags(d);
        REP(ai, i.size()) { bags[ai] += i[ai]; }
        REP(bi, j.size()) { bags[d - bi - 1] += j[bi]; }
        double variance = 0.0;
        for (const auto &k : bags) {
            variance += pow(k - ave_w, 2);
        }
        variance /= d;
        if (variance < ans) {
            ans = variance;
        }
    }
cout << fixed << setprecision(15) << ans << endl;

Python

for i in range(len(a)):
    a[i].sort(reverse=True)
for i in range(len(b)):
    b[i].sort(reverse=True)
ave_w = sum(w) / d
for i in a:
    for j in b:
        bags = [0] * d
        for ai, x in enumerate(i):
            bags[ai] += x
        for bi, y in enumerate(j):
            bags[~bi] += y
        variance = 0
        for k in bags:
            variance += (k - ave_w) ** 2
        ans = min(ans, variance / d)
print(ans)

全体の実装

C++

クリックで表示

#include <bits/stdc++.h>
using namespace std;
#include <atcoder/all>
using namespace atcoder;
typedef long long ll;
#define REP(i, n) for (ll i = 0; i < (ll)(n); i++)
int main() {
    cin.tie(nullptr);
    ios_base::sync_with_stdio(false);
    ll n;
    cin >> n;
    ll d;
    cin >> d;
    vector<ll> w(n);
    REP(i, n) { cin >> w[i]; }
    vector<vector<ll>> a, b;
    vector<ll> v;
    auto dfs = [&](auto self, vector<vector<ll>> &vec, ll start, ll end,
                   ll index) -> void {
        if (index == end) {
            vec.push_back(v);
            sort(vec.back().rbegin(), vec.back().rend());
            return;
        }
        REP(i, v.size()) {
            v[i] += w[index];
            self(self, vec, start, end, index + 1);
            v[i] -= w[index];
        }
        if (ll(v.size()) < d) {
            v.push_back(w[index]);
            self(self, vec, start, end, index + 1);
            v.pop_back();
        }
    };
    dfs(dfs, a, 0, n / 2, 0);
    dfs(dfs, b, n / 2, n, n / 2);
    for (auto &&i : a) {
        sort(i.rbegin(), i.rend());
    }
    for (auto &&i : b) {
        sort(i.rbegin(), i.rend());
    }
    double ans = 1e18;
    double ave_w = accumulate(w.begin(), w.end(), 0.0) / d;
    for (const auto &i : a)
        for (const auto &j : b) {
            vector<ll> bags(d);
            REP(ai, i.size()) { bags[ai] += i[ai]; }
            REP(bi, j.size()) { bags[d - bi - 1] += j[bi]; }
            double variance = 0.0;
            for (const auto &k : bags) {
                variance += pow(k - ave_w, 2);
            }
            variance /= d;
            if (variance < ans) {
                ans = variance;
            }
        }
    cout << fixed << setprecision(15) << ans << endl;
}

Python

クリックで表示

def resolve():
    import sys

    input = sys.stdin.readline
    n, d = map(int, input().split())
    w = tuple(map(int, input().split()))
    ans = float("inf")

    def divide_team(start, end):
        q = [(start, [])]
        while q:
            i, a = q.pop()
            if i == end:
                yield a
                continue
            for j in range(len(a)):
                b = a[:]
                b[j] += w[i]
                q.append((i + 1, b))
            if len(a) < d:
                q.append((i + 1, a + [w[i]]))

    a = tuple(divide_team(0, n // 2))
    b = tuple(divide_team(n // 2, n))
    for i in range(len(a)):
        a[i].sort(reverse=True)
    for i in range(len(b)):
        b[i].sort(reverse=True)
    ave_w = sum(w) / d
    for i in a:
        for j in b:
            bags = [0] * d
            for ai, x in enumerate(i):
                bags[ai] += x
            for bi, y in enumerate(j):
                bags[~bi] += y
            variance = 0
            for k in bags:
                variance += (k - ave_w) ** 2
            ans = min(ans, variance / d)
    print(ans)


if __name__ == "__main__":
    resolve()