Problem 1004: Balanced Integer

Problem 1004 source code

C++

#include <array>
#include <cassert>
#include <cstdint>
#include <iostream>
#include <vector>

namespace {

using i64 = std::int64_t;

constexpr i64 MOD = 1'000'000'007LL;
constexpr int DIGITS = 10;
constexpr int MAX_CELLS = DIGITS * DIGITS;

std::array<i64, MAX_CELLS + 1> factorial;

i64 mod_pow(i64 base, i64 exp) {
    i64 result = 1;
    while (exp > 0) {
        if ((exp & 1LL) != 0) {
            result = result * base % MOD;
        }
        base = base * base % MOD;
        exp >>= 1LL;
    }
    return result;
}

i64 mod_inverse(const i64 value) {
    return mod_pow(value, MOD - 2);
}

i64 shape_word_count(const std::vector<int>& partition) {
    int cells = 0;
    for (const int row : partition) {
        cells += row;
    }

    i64 hook_product = 1;
    i64 content_product = 1;
    for (int i = 0; i < static_cast<int>(partition.size()); ++i) {
        for (int j = 0; j < partition[i]; ++j) {
            int below = 0;
            for (int r = i + 1; r < static_cast<int>(partition.size()); ++r) {
                if (partition[r] > j) {
                    ++below;
                }
            }

            const int right = partition[i] - j - 1;
            const int hook = right + below + 1;
            const int content = DIGITS + (j + 1) - (i + 1);
            hook_product = hook_product * hook % MOD;
            content_product = content_product * content % MOD;
        }
    }

    const i64 inv_hooks = mod_inverse(hook_product);
    return factorial[cells] * content_product % MOD * inv_hooks % MOD * inv_hooks % MOD;
}

void enumerate_partitions(const int max_part,
                          const int max_rows,
                          std::vector<int>& partition,
                          const int max_cells,
                          i64& balanced,
                          i64& decreasing_excess) {
    if (!partition.empty()) {
        int cells = 0;
        for (const int row : partition) {
            cells += row;
        }

        if (cells <= max_cells) {
            const i64 ways = shape_word_count(partition);
            const int width = partition.front();
            const int height = static_cast<int>(partition.size());
            if (width == height) {
                balanced += ways;
                if (balanced >= MOD) {
                    balanced -= MOD;
                }
            }
            if (height == width + 1) {
                decreasing_excess += ways;
                if (decreasing_excess >= MOD) {
                    decreasing_excess -= MOD;
                }
            }
        }
    }

    if (static_cast<int>(partition.size()) == max_rows) {
        return;
    }

    int used = 0;
    for (const int row : partition) {
        used += row;
    }

    for (int next = max_part; next >= 1; --next) {
        if (used + next > max_cells) {
            continue;
        }
        partition.push_back(next);
        enumerate_partitions(next, max_rows, partition, max_cells, balanced, decreasing_excess);
        partition.pop_back();
    }
}

std::pair<i64, i64> count_all_words(const int max_cells) {
    i64 balanced = 0;
    i64 decreasing_excess = 0;
    std::vector<int> partition;
    enumerate_partitions(DIGITS, DIGITS, partition, max_cells, balanced, decreasing_excess);
    return {balanced, decreasing_excess};
}

i64 positive_balanced_count(const int max_digits) {
    const auto [balanced, ignored] = count_all_words(max_digits);
    const auto [unused, decreasing_excess] = count_all_words(max_digits - 1);
    (void)ignored;
    (void)unused;

    i64 result = (balanced - decreasing_excess - 1) % MOD;
    if (result < 0) {
        result += MOD;
    }
    return result;
}

void run_checkpoints() {
    assert(positive_balanced_count(4) == 2274);
}

}  // namespace

int main() {
    factorial[0] = 1;
    for (int i = 1; i <= MAX_CELLS; ++i) {
        factorial[i] = factorial[i - 1] * i % MOD;
    }

    run_checkpoints();
    std::cout << positive_balanced_count(MAX_CELLS) << '\n';
    return 0;
}

Python

#!/usr/bin/env python3
"""Project Euler Problem 1004 - Balanced Integer."""

from __future__ import annotations

MOD = 1_000_000_007
DIGITS = 10
MAX_CELLS = DIGITS * DIGITS

factorial = [1] * (MAX_CELLS + 1)
for i in range(1, MAX_CELLS + 1):
    factorial[i] = factorial[i - 1] * i % MOD


def mod_inverse(value: int) -> int:
    return pow(value, MOD - 2, MOD)


def shape_word_count(partition: list[int]) -> int:
    cells = sum(partition)
    hook_product = 1
    content_product = 1

    for i, row_len in enumerate(partition):
        for j in range(row_len):
            below = sum(1 for later_row in partition[i + 1 :] if later_row > j)
            right = row_len - j - 1
            hook = right + below + 1
            content = DIGITS + (j + 1) - (i + 1)
            hook_product = hook_product * hook % MOD
            content_product = content_product * content % MOD

    inv_hooks = mod_inverse(hook_product)
    return factorial[cells] * content_product % MOD * inv_hooks % MOD * inv_hooks % MOD


def enumerate_partitions(
    max_part: int,
    max_rows: int,
    partition: list[int],
    max_cells: int,
) -> tuple[int, int]:
    balanced = 0
    decreasing_excess = 0

    if partition:
        cells = sum(partition)
        if cells <= max_cells:
            ways = shape_word_count(partition)
            width = partition[0]
            height = len(partition)
            if width == height:
                balanced = (balanced + ways) % MOD
            if height == width + 1:
                decreasing_excess = (decreasing_excess + ways) % MOD

    if len(partition) == max_rows:
        return balanced, decreasing_excess

    used = sum(partition)
    for next_part in range(max_part, 0, -1):
        if used + next_part > max_cells:
            continue
        partition.append(next_part)
        sub_balanced, sub_excess = enumerate_partitions(
            next_part, max_rows, partition, max_cells
        )
        balanced = (balanced + sub_balanced) % MOD
        decreasing_excess = (decreasing_excess + sub_excess) % MOD
        partition.pop()

    return balanced, decreasing_excess


def count_all_words(max_cells: int) -> tuple[int, int]:
    return enumerate_partitions(DIGITS, DIGITS, [], max_cells)


def positive_balanced_count(max_digits: int) -> int:
    balanced, _ = count_all_words(max_digits)
    _, decreasing_excess = count_all_words(max_digits - 1)
    return (balanced - decreasing_excess - 1) % MOD


def run_checkpoints() -> None:
    assert positive_balanced_count(4) == 2274


def main() -> None:
    run_checkpoints()
    print(positive_balanced_count(MAX_CELLS))


if __name__ == "__main__":
    main()

Java

import java.util.ArrayList;
import java.util.List;

public class Euler1004 {
    private static final long MOD = 1_000_000_007L;
    private static final int DIGITS = 10;
    private static final int MAX_CELLS = DIGITS * DIGITS;
    private static final long[] FACTORIAL = new long[MAX_CELLS + 1];

    private static final class Counts {
        long balanced;
        long decreasingExcess;
    }

    private static long modPow(long base, long exp) {
        long result = 1;
        while (exp > 0) {
            if ((exp & 1L) != 0) {
                result = result * base % MOD;
            }
            base = base * base % MOD;
            exp >>= 1;
        }
        return result;
    }

    private static long modInverse(long value) {
        return modPow(value, MOD - 2);
    }

    private static long shapeWordCount(List<Integer> partition) {
        int cells = 0;
        for (int row : partition) {
            cells += row;
        }

        long hookProduct = 1;
        long contentProduct = 1;
        for (int i = 0; i < partition.size(); ++i) {
            int rowLength = partition.get(i);
            for (int j = 0; j < rowLength; ++j) {
                int below = 0;
                for (int r = i + 1; r < partition.size(); ++r) {
                    if (partition.get(r) > j) {
                        ++below;
                    }
                }

                int right = rowLength - j - 1;
                int hook = right + below + 1;
                int content = DIGITS + (j + 1) - (i + 1);
                hookProduct = hookProduct * hook % MOD;
                contentProduct = contentProduct * content % MOD;
            }
        }

        long invHooks = modInverse(hookProduct);
        return FACTORIAL[cells] * contentProduct % MOD * invHooks % MOD * invHooks % MOD;
    }

    private static Counts enumeratePartitions(
            int maxPart,
            int maxRows,
            List<Integer> partition,
            int maxCells) {
        Counts result = new Counts();

        if (!partition.isEmpty()) {
            int cells = 0;
            for (int row : partition) {
                cells += row;
            }

            if (cells <= maxCells) {
                long ways = shapeWordCount(partition);
                int width = partition.get(0);
                int height = partition.size();
                if (width == height) {
                    result.balanced = (result.balanced + ways) % MOD;
                }
                if (height == width + 1) {
                    result.decreasingExcess = (result.decreasingExcess + ways) % MOD;
                }
            }
        }

        if (partition.size() == maxRows) {
            return result;
        }

        int used = 0;
        for (int row : partition) {
            used += row;
        }

        for (int next = maxPart; next >= 1; --next) {
            if (used + next > maxCells) {
                continue;
            }
            partition.add(next);
            Counts child = enumeratePartitions(next, maxRows, partition, maxCells);
            result.balanced = (result.balanced + child.balanced) % MOD;
            result.decreasingExcess = (result.decreasingExcess + child.decreasingExcess) % MOD;
            partition.remove(partition.size() - 1);
        }

        return result;
    }

    private static Counts countAllWords(int maxCells) {
        return enumeratePartitions(DIGITS, DIGITS, new ArrayList<>(), maxCells);
    }

    private static long positiveBalancedCount(int maxDigits) {
        Counts all = countAllWords(maxDigits);
        Counts leadingZeroTails = countAllWords(maxDigits - 1);
        long result = (all.balanced - leadingZeroTails.decreasingExcess - 1) % MOD;
        if (result < 0) {
            result += MOD;
        }
        return result;
    }

    private static void runCheckpoints() {
        if (positiveBalancedCount(4) != 2274) {
            throw new AssertionError("checkpoint failed");
        }
    }

    public static void main(String[] args) {
        FACTORIAL[0] = 1;
        for (int i = 1; i <= MAX_CELLS; ++i) {
            FACTORIAL[i] = FACTORIAL[i - 1] * i % MOD;
        }

        runCheckpoints();
        System.out.println(positiveBalancedCount(MAX_CELLS));
    }
}