Wednesday, August 24, 2016

Codeforces Round #365 (Div. 2) - E. Mishka and Divisors

Problem Statement:
http://codeforces.com/contest/703/problem/E

Summary:
Given an array of n <= 1000 integers a[i] <= 10^12, and an integer k <= 10^12. Find subset of the array which product of the elements is divisible by k, such that the size of the subset is minimised. If more than one such subsets exist, return the one with smallest sum over the elements.


Solution:
The tricky part is to define the DP state to keep track. Suppose we solve a much simpler form of the problem: compute the subset of minimum size such that the product is divisible by k. Let dp[i][d] be the subset of a[1..i] with minimum size (minimum subsequence length) which product is divisible by d. For d0 | d1 (d1 is divisible by d0), it is true that dp[i][d1] >= dp[i][d0]. (Based on the observation that if S is the minimum subset S which product is divisible by d1, then S will also be divisible by d0, which means that length of S is the upper limit for dp[i][d0].)

To compute dp[i][d], we have two choices:
1. Pick a[i] into our subset. This gives us dp[i-1][d/gcd(d, a[i])] + 1. In other words, the minimum length if we pick a[i] is achieved by finding the subset in [1..i-1] that is divisible by the remaining factor of d that is not covered by a[i]. This is the key relationship in the DP transition that allows us to solve the problem. (Correctness argument: If we pick a[i], then we need to find the minimum of dp[i-1][X] for all X such that a[i]*X is divisible by d. The smallest of such X is exactly d/gcd(a[i], d), which also corresponds to the desired minimum.)
2. Don't pick a[i], then the minimum subset possible is derived from dp[i-1][d].
Hence dp[i][d] = min(dp[i-1][d/gcd(d, a[i])] + 1, d[i-1][d]).

To reduce the dp state, we can instead compute the prime factors of k, and hence notice that what matters is to keep track of the prime factors in each of a[i]. As it turns out, the number of divisors possible given k is still in the order of 10^3. I think one can check for this by enumerating over all possible prime factors and multiplicity, maximising the number of prime factors by choosing smaller primes.

Implementation:

#include <cstdio>
#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;
vector<int> primes;
int prime_marker[1000000];
void PrecomputePrimes() {
  const int MAX_PRIME = 1000000;
  for (int i = 2; i * i < MAX_PRIME; ++i) {
    for (int j = i; i * j < MAX_PRIME; ++j) {
      prime_marker[i*j]=1;
    }
  }
  for (int i = 2; i < MAX_PRIME; ++i) {
    if (prime_marker[i]) continue;
    primes.push_back(i);
  }
}

// Prime factors of K.
vector<int64_t> prime_factors;
vector<int> factor_max_power;
void ComputePrimeFactors(int64_t val) {
  for (int i = 0; i < primes.size(); ++i) {
    if (val % primes[i] == 0) {
      prime_factors.push_back(primes[i]);
      factor_max_power.push_back(0);
      while (val % primes[i] == 0) {
        val /= primes[i];
        factor_max_power.back()++;
      }
    }
  }
  if (val > 1000000) {
    prime_factors.push_back(val);
    factor_max_power.push_back(1);
  }
}

vector<int64_t> arr;
vector<int> cofactor_power;
void ComputeCofactorPower(int64_t val) {
  cofactor_power.resize(prime_factors.size());
  for (int i = 0; i < prime_factors.size(); ++i) {
    cofactor_power[i] = 0;
    if (val % prime_factors[i] == 0) {
      while (val % prime_factors[i] == 0) {
        val /= prime_factors[i];
        cofactor_power[i]++;
      }
    }
  }
}

//       Indices            Hash
// E.g. [i][j][k] = i * NJ * NK + j * NK + k.
vector<int> basis;
void InitializeBasis() {
  basis.resize(prime_factors.size());
  basis[0] = 1;
  for (int i = 1; i < prime_factors.size(); ++i) {
    basis[i] = basis[i-1] * (factor_max_power[i-1]+1);
  }
}
void MapToIndices(int hash, vector<int>* indices /* OUT */) {
  for (int i = basis.size()-1; i >= 0; --i) {
    (*indices)[i] = hash / basis[i];
    hash -= (*indices)[i] * basis[i];
  }
}
int IndicesToHash(const vector<int>& indices) {
  int hash = 0;
  for (int i = 0; i < indices.size(); ++i) {
    hash += indices[i]*basis[i];
  }
  return hash;
}

vector<vector<int>> hash_indices;
void InitHashIndices(int max_hash) {
  hash_indices.resize(max_hash, vector<int>(basis.size()));
  for (int i = 0; i < max_hash; ++i) {
    MapToIndices(i, &hash_indices[i]);
  }
}

struct DpState {
  int length = 100000;
  int64_t cost = 0;
  int parent;
  bool is_chosen = false;

  DpState(int length, int64_t cost, int parent, bool is_chosen)
      : length(length), cost(cost), parent(parent), is_chosen(is_chosen) {}
  DpState() {}

  bool operator<(const DpState& other) const {
    if (length == other.length) return cost < other.cost;
    return length < other.length;
  }
};

DpState dp[1003][7000];
void UpdateDP(int cur, int hash, const vector<int>& indices) {  
  if (cur == 0) {
    if (hash == 0) {
      dp[cur][hash] = DpState{0, 0, -1, false};
      return;
    }
    bool is_divisible = true;
    for (int i = 0; i < indices.size(); ++i) {
      if (indices[i] > cofactor_power[i]) {
        is_divisible = false;
        break;
      }
    }
    if (is_divisible) {
      dp[cur][hash] = DpState{
        1 /* length */,
        arr[cur] /* cost */,
        -1 /* parent */,
        true
      };
    }
    return;
  }
  // Case 1: Take arr[cur].
  vector<int> prev_indices(indices.size());
  for (int i = 0; i < indices.size(); ++i) {
    prev_indices[i] = max(0, indices[i] - cofactor_power[i]);
  }
  int prev_hash = IndicesToHash(prev_indices);
  DpState state1 = {
    dp[cur-1][prev_hash].length + 1,
    dp[cur-1][prev_hash].cost + arr[cur],
    prev_hash,
    true
  };
  // Case 2: Don't take arr[cur].
  DpState state2 = {
    dp[cur-1][hash].length,
    dp[cur-1][hash].cost,
    hash,
    false
  };
  if (state1 < state2) {
    dp[cur][hash] = state1;
  } else {
    dp[cur][hash] = state2;
  }
}

int main(){
  int n;
  int64_t k;
  PrecomputePrimes();
  scanf("%d",&n);
  cin >> k;
  int64_t val;
  for (int i = 0; i < n; ++i) {
    cin >> val;
    arr.push_back(val);
  }
  // Special Case: d = 1.
  if (k == 1) {
    int lowest_index = 0;
    for (int i = 1; i < n; ++i) {
      if (arr[i] < arr[lowest_index]) lowest_index = i;
    }
    printf("1\n%d\n",lowest_index+1);
    return 0;
  }
  ComputePrimeFactors(k);
  InitializeBasis();
  int max_hash = basis.back() * (factor_max_power.back()+1);
  InitHashIndices(max_hash);
  for (int i = 0; i < n; ++i) {
    ComputeCofactorPower(arr[i]);
    for (int hash = 0; hash < max_hash; ++hash) {
      UpdateDP(i, hash, hash_indices[hash]);
    }
  }
  int cur_hash = max_hash-1;
  if (dp[n-1][cur_hash].length >= 100000) {
    printf("-1\n");
    return 0;
  }
  printf("%d\n", dp[n-1][cur_hash].length);
  for (int i = n-1; i>=0; --i) {
    if(dp[i][cur_hash].is_chosen) {
      printf("%d ", i+1);
    }
    cur_hash = dp[i][cur_hash].parent;
  }
  printf("\n");
  return 0;
}

The time complexity is hence O(n * number of primes * number of possible divisors).