Submission #636708

#	Time	Username	Problem	Language	Result	Execution time	Memory
636708		abeker	Žarulje (COI15_zarulje)	C++17	100 / 100	348 ms	25196 KiB

zarulje

#include <bits/stdc++.h>
using namespace std;

const int MAXN = 2e5 + 5;
const int MOD = 1e9 + 7;

int N, K;
vector <int> pos[MAXN];
int fact[MAXN], inv[MAXN];
int equal_lft[MAXN], equal_rig[MAXN];
int ans[MAXN];

void load() {
  scanf("%d%d", &N, &K);
  for (int i = 1; i <= N; i++) {
    int x;
    scanf("%d", &x);
    pos[x].push_back(i);
  }
}

inline int mul(int x, int y) {
  return (long long)x * y % MOD;
}

inline void mult(int &x, int y) {
  x = mul(x, y);
}

int pot(int x, int y) {
  int res = 1;
  for (; y; y /= 2) {
    if (y % 2)
      res = mul(res, x);
    x = mul(x, x);
  }
  return res;
}

void inc(int lo, int hi, int a, int b) {
  mult(ans[lo], mul(fact[a + b], mul(inv[a], inv[b])));
  mult(ans[hi], mul(mul(fact[a], fact[b]), inv[a + b]));
}

void solve() {
  fact[0] = inv[0] = 1;
  for (int i = 1; i <= N; i++) {
    fact[i] = mul(fact[i - 1], i);
    inv[i] = pot(fact[i], MOD - 2);
    ans[i] = 1;
  }
  set <int> sofar;
  for (int i = 1; i < MAXN; i++) {
    auto get_prev = [&](int x) {
      if (sofar.empty())
        return 0;
      auto it = sofar.lower_bound(x);
      return it == sofar.begin() ? 0 : *--it;
    };
    auto get_next = [&](int x) {
      if (sofar.empty())
        return N + 1;
      auto it = sofar.lower_bound(x + 1);
      return it == sofar.end() ? N + 1 : *it;
    };
    for (auto it : pos[i]) {
      auto count_equal = [&](int lo, int hi) {
        return lower_bound(pos[i].begin(), pos[i].end(), hi) - lower_bound(pos[i].begin(), pos[i].end(), lo);
      };
      equal_lft[it] = count_equal(get_prev(it) + 1, it + 1);
      equal_rig[it] = count_equal(it, get_next(it));
    }
    for (auto it : pos[i])
      sofar.insert(it);
    for (int j = 0; j < (int)pos[i].size() - 1; j++) {
      auto attempt = [&](int l, int r) {
        int tmp = get_next(l);
        if (tmp >= r)
          inc(l + 1, r, equal_lft[l], equal_rig[r]);
        else if (get_next(tmp) >= r)
          inc(tmp, tmp + 1, equal_lft[l], equal_rig[r]);
      };
      attempt(pos[i][j], pos[i][j + 1]);
      if (j < (int)pos[i].size() - 2)
        attempt(pos[i][j], pos[i][j + 2]);
    }
  }
  for (int i = 2; i <= N; i++)
    mult(ans[i], ans[i - 1]);
  while (K--) {
    int p;
    scanf("%d", &p);
    printf("%d ", ans[p]);
  }
  puts("");
}

int main() {
  load();
  solve();
  return 0;
}

Compilation message (stderr)

zarulje.cpp: In function 'void load()':
zarulje.cpp:14:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   14 |   scanf("%d%d", &N, &K);
      |   ~~~~~^~~~~~~~~~~~~~~~
zarulje.cpp:17:10: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   17 |     scanf("%d", &x);
      |     ~~~~~^~~~~~~~~~
zarulje.cpp: In function 'void solve()':
zarulje.cpp:92:10: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   92 |     scanf("%d", &p);
      |     ~~~~~^~~~~~~~~~

• Subtask 1: 22 / 22
• Subtask 2: 38 / 38
• Subtask 3: 40 / 40