subset_operations.cpp

#ifndef SUBSETS_OPERATIONS
#define SUBSETS_OPERATIONS

// The class SubsetFunction<T> represents a function on the subsets of
// {0, 1, ..., N-1} (or equivalently the integers [0, 2^N)) with values in T
// and offers a variety of standard operations (i.e. xor-convolution).
//
// The time complexities of the various operations are optimal.
// Even if not heavily optimized (in order to maintain a good readability),
// these implementations shall be fast enough in all situations.
//
// Only the high-level APIs (that is SubsetFunction<T>) shall be used by the
// end-user. There is no reason to call directly the low-level functions in
// namespace::internal (which, for performance and for simplicity of the
// implementation, are heavily using pointers).

#include <cassert>
#include <algorithm>
using namespace std;

// In namespace::internal, all operations are implemented directly on arrays
// (and, whenever possible, in place). There are no memory leaks.
// Operating directly on arrays simplifies the implementation (since often one
// has to treat a subarray as an array and with vectors this is not possible).
//
// The functions in namespace::internal operate on pointers and are, therefore, a bit
// uncomfortable to use for the end-user. It is much better to use the clean
// APIs offered by the class SubsetFunction.
namespace internal {

// All functions in namespace::internal take as first argument N, and as further
// arguments some arrays with size 2^N.
// Each function accepts exactly one non-const argument, that argument will
// contain (after the execution) the result of the function.

// Reset all values to 0.
// Complexity: O(2^N).
template<typename T>
void clear(int N, T* A) { fill(A, A+(1<<N), 0); }

// A'[x] = A[complement of x]
// Complexity: O(2^N).
template<typename T>
void complement(int N, T* A) {
    int pot = 1<<N;
    for (int x = 0; x < (pot>>1); x++) swap(A[x], A[x^(pot-1)]);
}

// In place xor-transform. See xor_convolution to understand its importance.
// A'[x] = \sum_y A[y]*(-1)^{bits(x&y)}.
// Complexity: O(N*2^N).
template<typename T>
void xor_transform(int N, T* A, bool inverse = false) {
    for (int len = 1; 2 * len <= (1<<N); len *= 2) {
        for (int x = 0; x < (1<<N); x += 2 * len) {
            for (int y = 0; y < len; y++) {
                T u = A[x + y];
                T v = A[x + y + len];
                A[x + y] = u + v;
                A[x + y + len] = u - v;
            }
        }
    }
    if (inverse) for (int i = 0; i < (1<<N); i++) A[i] /= (1<<N);
}

// In place and-transform. See and_convolution to understand its importance.
//   A'[x] = \sum_{x\subseteq y} A[y].
// Complexity: O(N*2^N).
template<typename T>
void and_transform(int N, T* A, bool inverse = false) {
    for (int len = 1; 2 * len <= (1<<N); len *= 2) {
        for (int x = 0; x < (1<<N); x += 2 * len) {
            for (int y = 0; y < len; y++) {
                if (inverse) A[x+y] -= A[x+y+len];
                else A[x+y] += A[x+y+len];
            }
        }
    }
}

// In place or-transform. See or_convolution to understand its importance.
//   A'[x] = \sum_{y\subseteq x} A[y].
// Complexity: O(N*2^N).
template<typename T>
void or_transform(int N, T* A, bool inverse = false) {
    for (int len = 1; 2 * len <= (1<<N); len *= 2) {
        for (int x = 0; x < (1<<N); x += 2 * len) {
            for (int y = 0; y < len; y++) {
                if (inverse) A[x+y+len] -= A[x+y];
                else A[x+y+len] += A[x+y];
            }
        }
    }
}


// C[x] = \sum_{y^z = x} A[y]B[z]
// Complexity: O(N2^N).
template<typename T>
void xor_convolution(int N, const T* A, const T* B, T* C) {
    T* B2 = new T[1<<N];
    for (int x = 0; x < (1<<N); x++) C[x] = A[x], B2[x] = B[x];
    xor_transform(N, C);
    xor_transform(N, B2);
    for (int x = 0; x < (1<<N); x++) C[x] *= B2[x];
    xor_transform(N, C, true);
    delete[] B2;
}

// C[x] = \sum_{y&z = x} A[y]B[z]
// Complexity: O(N2^N).
template<typename T>
void and_convolution(int N, const T* A, const T* B, T* C) {
    T* B2 = new T[1<<N];
    for (int x = 0; x < (1<<N); x++) C[x] = A[x], B2[x] = B[x];
    and_transform(N, C);
    and_transform(N, B2);
    for (int x = 0; x < (1<<N); x++) C[x] *= B2[x];
    and_transform(N, C, true);
    delete[] B2;
}

// C[x] = \sum_{y|z = x} A[y]B[z]
// Complexity: O(N2^N).
template<typename T>
void or_convolution(int N, const T* A, const T* B, T* C) {
    T* B2 = new T[1<<N];
    for (int x = 0; x < (1<<N); x++) C[x] = A[x], B2[x] = B[x];
    or_transform(N, C);
    or_transform(N, B2);
    for (int x = 0; x < (1<<N); x++) C[x] *= B2[x];
    or_transform(N, C, true);
    delete[] B2;
}


// C[x] = \sum_{y&z = 0, y|z = x} A[y]*B[z].
// Complexity: O(N²*2^N).
template<typename T>
void subset_convolution(int N, const T* A, const T* B, T* C) {
    T** A_popcnt = new T*[N+1];
    T** B_popcnt = new T*[N+1];
    T* tmp = new T[1<<N];
    for (int i = 0; i <= N; i++) {
        A_popcnt[i] = new T[1<<N];
        B_popcnt[i] = new T[1<<N];
        clear(N, A_popcnt[i]);
        clear(N, B_popcnt[i]);
    }
    for (int x = 0; x < (1<<N); x++) {
        int q = __builtin_popcount(x);
        A_popcnt[q][x] = A[x];
        B_popcnt[q][x] = B[x];
    }
    for (int i = 0; i <= N; i++) {
        or_transform(N, A_popcnt[i]);
        or_transform(N, B_popcnt[i]);
    }
    for (int l = 0; l <= N; l++) {
        clear(N, tmp);
        for (int i = 0; i <= l; i++) {
            int j = l-i;
            for (int x = 0; x < (1<<N); x++)
                tmp[x] += A_popcnt[i][x]*B_popcnt[j][x];
        }
        or_transform(N, tmp, true);
        for (int x = 0; x < (1<<N); x++)
            if (__builtin_popcount(x) == l) C[x] = tmp[x];
    }
    for (int i = 0; i <= N; i++) {
        delete[] A_popcnt[i];
        delete[] B_popcnt[i];
    }
    delete[] A_popcnt;
    delete[] B_popcnt;
    delete[] tmp;
}

// subset_convolution(A, C) = B.
// Complexity: O(N²2^N).
//
// ACHTUNG: The value A[0] must be invertible.
template<typename T>
void inverse_subset_convolution(int N, const T* A, const T* B, T* C) {
    T inv = 1/A[0];
    
    T** A_popcnt = new T*[N+1];
    T** C_popcnt = new T*[N+1];
    for (int i = 0; i <= N; i++) {
        A_popcnt[i] = new T[1<<N];
        C_popcnt[i] = new T[1<<N];
        clear(N, A_popcnt[i]);
        clear(N, C_popcnt[i]);
    }
    for (int x = 0; x < (1<<N); x++) {
        int q = __builtin_popcount(x);
        A_popcnt[q][x] = A[x];
    }
    for (int i = 0; i <= N; i++) or_transform(N, A_popcnt[i]);
    
    for (int l = 0; l <= N; l++) {
        for (int i = 0; i <= l; i++) {
            int j = l-i;
            for (int x = 0; x < (1<<N); x++)
                C_popcnt[l][x] += A_popcnt[i][x]*C_popcnt[j][x];
        }
        or_transform(N, C_popcnt[l], true);
        for (int x = 0; x < (1<<N); x++) {
            int q = __builtin_popcount(x);
            if (q != l) C_popcnt[l][x] = 0;
            else {
                C_popcnt[l][x] = (B[x] - C_popcnt[l][x])*inv;
                C[x] = C_popcnt[l][x];
            }
        }
        or_transform(N, C_popcnt[l]);
    }

    for (int i = 0; i <= N; i++) {
        delete[] A_popcnt[i];
        delete[] C_popcnt[i];
    }
    delete[] A_popcnt;
    delete[] C_popcnt;
}


// B = exp(A).
// Complexity: O(N²2^N).
//
// For x != 0, it holds
// B[x] = \sum_{0 < y_1 < y_2 < ... < y_k: y_i disjoint, y_1|y_2|...|y_k = x}
//                    A[y_1]*A[y_2]*...*A[y_k].
// The value B[0] is (arbitrarily) set to 1.
//
// This operations is denoted exponential in analogy with the exponential of
// classical power-series. Indeed, a function on the subsets might also be
// interpreted as power series: the exponent is a subset and the coefficient
// is the values of the function for such subset.
// With this interpretation in mind, assuming that the multiplication is given
// by the subset_convolution, there is a pretty clear analogy between this
// operation and the classical exponential.
//
// ACHTUNG: The value of A[0] is ignored.
template<typename T>
void exp(int N, const T* A, T* B) {
    B[0] = 1;
    for (int n = 0; n < N; n++) 
        subset_convolution(n, A + (1<<n), B, B + (1<<n));
}

// exp(B) = A.
// Complexity: O(N²2^N).
//
// The value of B[0] is (arbitrarily) set to 0.
//
// ACHTUNG: It must hold A[0] = 1.
template<typename T>
void log(int N, const T* A, T* B) {
    assert(A[0] == 1);
    B[0] = 0;
    for (int n = 0; n < N; n++)
        inverse_subset_convolution(n, A, A+(1<<n), B+(1<<n));
}

}; // namespace internal



// The class SubsetFunction<T> represents a function on the subsets of
// {0, 1, ..., N-1} with values in T.
// A subset is represented by its bitmask. If Subset<T> F is an instance, then
// the value of F on the subset x is F[x].
//
// With respect to copy and assignment, the class SubsetFunction behaves like
// a vector.
// 
// The following operations (see their headers in namespace::internal to know
// precisely what they do) are supported:
//  xor_transform                  // Complexity O(N2^N).
//  and_transform                  // Complexity O(N2^N).
//  or_transform                   // Complexity O(N2^N).
//
//  xor_convolution                // Complexity O(N2^N).
//  and_convolution                // Complexity O(N2^N).
//  or_convolution                 // Complexity O(N2^N).
//  subset_convolution             // Complexity O(N²2^N).
//  inverse_subset_convolution     // Complexity O(N²2^N).
//
//  complement                     // Complexity O(2^N).
//  exp                            // Complexity O(N²2^N).
//  log                            // Complexity O(N²2^N).
//
// The code that follows is mostly boiler-plate, as the heavy-lifting is
// performed in namespace::internal.
//
// Usage Example:
//  SubsetFunction<double> A(10);
//  for (int x = 0; x < (1<<10); x++) cin >> A[i];
//  auto B = xor_convolution(A, A);
//  auto C = inverse_subset_convolution(A, B);
//  auto D = log(exp(C));
//  for (int x = 1; x < (1<<10); x++) assert(abs(D[x] - C[x]) < 0.0001);
//  auto E = and_transform(A);
//  E = and_transform(E, true); // inverse transform
//  for (int x = 0; x < (1<<10); x++) assert(abs(E[x] - A[x]) < 0.0001);
template<typename T>
struct SubsetFunction {
    int N; // 1<<N is the size of arr.
    T* arr;
    SubsetFunction(int N): N(N) { // N >= 0
        arr = new T[1<<N];
        clear();
    }
    ~SubsetFunction() { delete[] arr; }
    SubsetFunction(const SubsetFunction& other) {
        N = other.N;
        arr = new T[1<<N];
        for (int x = 0; x < (1<<N); x++) arr[x] = other[x];
    }
    SubsetFunction& operator=(const SubsetFunction& other) {
        if (this != &other) {
            if (N != other.N) {
                N = other.N;
                delete[] arr;
                arr = new T[1<<N];
            }
            for (int x = 0; x < (1<<N); x++) arr[x] = other[x];
        }
        return *this;
    }
    T& operator[](int index) { return arr[index]; }
    const T& operator[](int index) const { return arr[index]; }
    void clear() { internal::clear(N, arr); }
};

#define TRANSFORM_DEF(OP)                                         \
template<typename T>                                              \
SubsetFunction<T> OP##_transform(const SubsetFunction<T>& A,      \
                                 bool inverse=false) {            \
    SubsetFunction<T> B = A;                                      \
    internal::OP##_transform(B.N, B.arr, inverse);                \
    return B;                                                     \
}                                                                 \

TRANSFORM_DEF(xor)
TRANSFORM_DEF(and)
TRANSFORM_DEF(or)

#define CONVOLUTION_DEF(OP)                                         \
template<typename T>                                                \
SubsetFunction<T> OP##_convolution(const SubsetFunction<T>& A,      \
                                   const SubsetFunction<T>& B) {    \
    assert(A.N == B.N);                                             \
    SubsetFunction<T> C(A.N);                                       \
    internal::OP##_convolution(A.N, A.arr, B.arr, C.arr);           \
    return C;                                                       \
}

CONVOLUTION_DEF(xor)
CONVOLUTION_DEF(and)
CONVOLUTION_DEF(or)
CONVOLUTION_DEF(subset)
CONVOLUTION_DEF(inverse_subset)

#define UNARY_OPERATOR_DEF(OP)                          \
template<typename T>                                    \
SubsetFunction<T> OP(const SubsetFunction<T>& A) {      \
    SubsetFunction<T> B(A.N);                           \
    internal::OP(A.N, A.arr, B.arr);                    \
    return B;                                           \
}

UNARY_OPERATOR_DEF(exp)
UNARY_OPERATOR_DEF(log)

template<typename T>
SubsetFunction<T> complement(const SubsetFunction<T>& A) {
    SubsetFunction<T> B(A);
    internal::complement(B.N, B.arr);
    return B;
}


#endif // SUBSETS_OPERATIONS