diff --git a/arithmetic_progressions/arithmetic_progressions.c b/arithmetic_progressions/arithmetic_progressions.c index b9ee09e..8c48e91 100644 --- a/arithmetic_progressions/arithmetic_progressions.c +++ b/arithmetic_progressions/arithmetic_progressions.c @@ -16,6 +16,43 @@ struct adp { long p; }; +long pow_mod(long a, long b, long mod, long acc) { + long powa = a % mod; + while (b > 0) { + if (b & 1) { + acc = (acc * powa) % mod; + } + powa = (powa * powa) % mod; + b >>= 1; + } + return acc; +} + +long factorial_mod(long n, long mod, long acc) { + int i; + for (i = 2; i <= n; i++) { + acc = (acc * i) % mod; + } + return acc; +} + +long *fenwick_make(size_t n) { + long *fenwick = calloc(n, sizeof *fenwick); + return fenwick; +} + +long *fenwick_prodmake(size_t n) { + long *fenwick = malloc(n * sizeof *fenwick); + if (fenwick == NULL) { + return NULL; + } + size_t i; + for (i = 0; i < n; i++) { + fenwick[i] = 1; + } + return fenwick; +} + // Add val to element i in fenwick tree void fenwick_update(long *fenwick, size_t n, size_t i, long val) { i++; @@ -43,6 +80,35 @@ long fenwick_intervalsum(long *fenwick, size_t n, size_t i, size_t j) { return fenwick_prefixsum(fenwick, n, j) - fenwick_prefixsum(fenwick, n, i); } +// Multiply element i by val in fenwick tree of products +void fenwick_produpdate(long *fenwick, size_t n, size_t i, long val, long mod) { + i++; + if (i < 1 || i > n) { + return; + } + while (i <= n) { + fenwick[i-1] = (fenwick[i-1] * val) % mod; + i += i & -i; // parent + } +} + +// Product of elements in [0, i[ +long fenwick_prefixprod(long *fenwick, size_t n, size_t i, long mod) { + long prod = 1; + while (i > 0) { + prod = (prod * fenwick[i-1]) % mod; + i -= i & -i; // predecessor + } + return prod; +} + +// Product of elements in [i, j[ +long fenwick_intervalprod(long *fenwick, size_t n, size_t i, size_t j, long mod) { + long p_j = fenwick_prefixprod(fenwick, n, j, mod); + long p_i = fenwick_prefixprod(fenwick, n, i, mod); + // p_i ^ -1 = p_i ^ BIG_MOD - 2 in this field + return pow_mod(p_i, BIG_MOD - 2, BIG_MOD, p_j); +} // Add val to entire interval [i, j[ of the fenwick tree // Supposes the fenwick tree stores differences between elements void fenwick_intervalupdate( @@ -96,7 +162,7 @@ size_t read_ints(char *str, int **ints) { return i; } -int read_adp(int n, struct adp *adp, long *ptree) { +int read_adp(int n, struct adp *adp, long *ptree, long *vtree) { char buf[BUF_SIZE]; int *ints; size_t ints_len; @@ -111,64 +177,47 @@ int read_adp(int n, struct adp *adp, long *ptree) { adp[i].a = ints[0]; adp[i].d = ints[1]; adp[i].p = ints[2]; + long v = pow_mod(adp[i].d, adp[i].p, BIG_MOD, 1); + fenwick_produpdate(vtree, n, i, v, BIG_MOD); fenwick_update(ptree, n, i, ints[2]); free(ints); } return 0; } -long pow_mod(long a, long b, long mod, long acc) { - long powa = a % mod; - while (b > 0) { - if (b & 1) { - acc = (acc * powa) % mod; - } - powa = (powa * powa) % mod; - b >>= 1; - } - return acc; -} - -long factorial_mod(long n, long mod, long acc) { - int i; - for (i = 2; i <= n; i++) { - acc = (acc * i) % mod; - } - return acc; -} - -int min_const_diff(int n, struct adp *adp, long *ptree, int i, int j) { +int min_const_diff( + int n, + struct adp *adp, + long *ptree, + long *vtree, + int i, + int j) +{ if (i < 1 || j > n) { return 1; } - long v = 1; - long p; - int l; - for (l = i-1; l < j; l++) { - p = adp[l].p; - if (p > 0) { - v = pow_mod(adp[l].d, p, BIG_MOD, v); - } - } + long v = fenwick_intervalprod(vtree, n, i-1, j, BIG_MOD); long k = fenwick_intervalsum(ptree, n, i-1, j); v = factorial_mod(k, BIG_MOD, v); printf("%ld %ld\n", k, v); return 0; } -int add_powers(int n, struct adp *adp, long *ptree, int i, int j, int v) { +int add_powers(int n, struct adp *adp, long *ptree, long *vtree, int i, int j, int v) { if (i < 1 || j > n) { return 1; } int k; for (k = i-1; k < j; k++) { + long factor = pow_mod(adp[k].d, v, BIG_MOD, 1); + fenwick_produpdate(vtree, n, k, factor, BIG_MOD); fenwick_update(ptree, n, k, v); adp[k].p += v; } return 0; } -int handle_query(char *str, int n, struct adp *adp, long *ptree) { +int handle_query(char *str, int n, struct adp *adp, long *ptree, long *vtree) { int *ints; int ints_len = read_ints(str, &ints); if (ints_len < 1) { @@ -179,13 +228,13 @@ int handle_query(char *str, int n, struct adp *adp, long *ptree) { if (ints_len != 3) { return 1; } - return min_const_diff(n, adp, ptree, ints[1], ints[2]); + return min_const_diff(n, adp, ptree, vtree, ints[1], ints[2]); } if (query_type == 1) { if (ints_len != 4) { return 1; } - return add_powers(n, adp, ptree, ints[1], ints[2], ints[3]); + return add_powers(n, adp, ptree, vtree, ints[1], ints[2], ints[3]); } return 1; } @@ -201,11 +250,12 @@ int main(int argc, char **argv) { } struct adp *adp = malloc(n * sizeof *adp); - long *ptree = calloc(n, sizeof *ptree); - if (adp == NULL || ptree == NULL) { + long *ptree = fenwick_make(n); + long *vtree = fenwick_prodmake(n); + if (adp == NULL || ptree == NULL || vtree == NULL) { return 0; } - int err = read_adp(n, adp, ptree); + int err = read_adp(n, adp, ptree, vtree); if (err) { return 0; } @@ -219,7 +269,7 @@ int main(int argc, char **argv) { int i; for (i = 0; i < q; i++) { fgets(buf, BUF_SIZE, stdin); - err = handle_query(buf, n, adp, ptree); + err = handle_query(buf, n, adp, ptree, vtree); if (err) { return 0; }