Problem p02804 (C++)

Buggy submission (status: Time Limit Exceeded):
#define _GLIBCXX_DEBUG
#include <bits/stdc++.h>
#define For(i, a, b) for (int(i) = (a); (i) < (b); ++(i))
#define rFor(i, a, b) for (int(i) = (a)-1; (i) >= (b); --(i))
#define rep(i, n) For((i), 0, (n))
#define rrep(i, n) rFor((i), (n), 0)
#define fi first
#define se second
#define double long double
using namespace std;
typedef long long lint;
typedef unsigned long long ulint;
typedef pair<int, int> pii;
typedef pair<double, double> pdd;
typedef pair<lint, lint> pll;
typedef complex<double> xy_t;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

constexpr lint mod = 1e9 + 7;
constexpr lint INF = mod * mod;
constexpr int MAX = 100010;

struct mint {
  using i64 = int_fast64_t;
  i64 a;
  static constexpr i64 mod = 1e9 + 7;
  mint(const i64 a_ = 0) : a(a_) {}
  mint inv() {
    i64 t = 1, n = mod - 2, x = a;
    while (n) {
      if (n & 1)
        (t *= x) %= mod;
      (x *= x) %= mod;
      n >>= 1;
    }
    mint ret(t);
    return ret;
  }
  bool operator==(const mint &x) { return a == x.a; }
  bool operator!=(const mint &x) { return a != x.a; }
  mint operator+(const mint &x) {
    mint ret;
    ret.a = a + x.a;
    if (ret.a > mod)
      ret.a -= mod;
    return ret;
  }
  mint operator-(const mint &x) {
    mint ret;
    ret.a = a - x.a;
    if (ret.a < 0)
      ret.a += mod;
    return ret;
  }
  mint operator*(const mint &x) {
    mint ret;
    ret.a = a * x.a % mod;
    return ret;
  }
  mint operator/(mint &x) {
    mint ret;
    ret.a = a * x.inv().a % mod;
    return ret;
  }
  mint operator^(lint n) {
    mint ret = 1, x = a;
    while (n) {
      if (n & 1)
        (ret.a *= x.a) %= mod;
      (x.a *= x.a) %= mod;
      n >>= 1;
    }
    return ret;
  }
  mint &operator+=(const mint &x) {
    a += x.a;
    if (a >= mod)
      a -= mod;
    return *this;
  }
  mint &operator-=(const mint &x) {
    a -= x.a;
    if (a < 0)
      a += mod;
    return *this;
  }
  mint &operator*=(const mint &x) {
    (a *= x.a) %= mod;
    return *this;
  }
  mint &operator/=(mint &x) {
    (a *= x.inv().a) %= mod;
    return *this;
  }
  mint &operator^=(lint n) {
    lint ret = 1;
    while (n) {
      if (n & 1)
        (ret *= a) %= mod;
      (a *= a) %= mod;
      n >>= 1;
    }
    a = ret;
    return *this;
  }
};

vector<mint> fact;
vector<mint> revfact;

void setfact(int n) {
  fact.resize(n + 1);
  revfact.resize(n + 1);
  fact[0] = 1;
  rep(i, n) fact[i + 1] = fact[i] * mint(i + 1);

  revfact[n] = fact[n].inv();
  for (int i = n - 1; i >= 0; i--)
    revfact[i] = revfact[i + 1] * mint(i + 1);
}

mint getC(int n, int r) {
  if (n < r)
    return 0;
  return fact[n] * revfact[r] * revfact[n - r];
}

template <typename T> struct SegTree {
  using F = function<T(T, T)>;
  int sz = 1;
  T et;
  F f, g;
  vector<T> node;

  SegTree(int sz_, T et_, F f_, F g_) : et(et_), f(f_), g(g_) {
    while (sz < sz_)
      sz <<= 1;
    node.resize(sz << 1, et);
  }

  void update(int i, T x) {
    i += sz;
    node[i] = g(node[i], x);
    i >>= 1;
    while (i) {
      node[i] = f(node[i << 1], node[(i << 1) + 1]);
      i >>= 1;
    }
  }

  T query(int l, int r) {
    T vl = et, vr = et;
    for (l += sz, r += sz; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        vl = f(vl, node[l++]);
      if (r & 1)
        vr = f(node[--r], vr);
    }
    return f(vl, vr);
  }
};

int main() {
  int n, K;
  scanf("%d%d", &n, &K);
  int a[n];
  vector<int> v(n);
  rep(i, n) {
    scanf("%d", &a[i]);
    v[i] = a[i];
  }
  sort(v.begin(), v.end());
  v.erase(unique(v.begin(), v.end()), v.end());
  SegTree<int> st(v.size(), 0, plus<>(), plus<>());
  rep(i, n) {
    int idx = lower_bound(v.begin(), v.end(), a[i]) - v.begin();
    st.update(idx, 1);
  }
  setfact(n);
  mint ans = 0;
  sort(a, a + n);
  for (int t : v) {
    auto lb = lower_bound(a, a + n, t);
    auto ub = upper_bound(a, a + n, t);
    int x = lb - a;
    int y = ub - lb;
    int z = (a + n) - ub;
    auto P = (getC(y + z, K) - getC(z, K)) * mint((t + mod) % mod);
    auto Q = (getC(x + y, K) - getC(x, K)) * mint((t + mod) % mod);
    ans += Q - P;
    // printf("%lld %lld\n", P.a, Q.a);
    // printf("%d %d %d %d\n", x, y, z, t);
  }
  printf("%lld\n", ans.a);
}

Fixed submission:
#include <bits/stdc++.h>
#define For(i, a, b) for (int(i) = (a); (i) < (b); ++(i))
#define rFor(i, a, b) for (int(i) = (a)-1; (i) >= (b); --(i))
#define rep(i, n) For((i), 0, (n))
#define rrep(i, n) rFor((i), (n), 0)
#define fi first
#define se second
#define double long double
using namespace std;
typedef long long lint;
typedef unsigned long long ulint;
typedef pair<int, int> pii;
typedef pair<double, double> pdd;
typedef pair<lint, lint> pll;
typedef complex<double> xy_t;
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

constexpr lint mod = 1e9 + 7;
constexpr lint INF = mod * mod;
constexpr int MAX = 100010;

struct mint {
  using i64 = int_fast64_t;
  i64 a;
  static constexpr i64 mod = 1e9 + 7;
  mint(const i64 a_ = 0) : a(a_) {}
  mint inv() {
    i64 t = 1, n = mod - 2, x = a;
    while (n) {
      if (n & 1)
        (t *= x) %= mod;
      (x *= x) %= mod;
      n >>= 1;
    }
    mint ret(t);
    return ret;
  }
  bool operator==(const mint &x) { return a == x.a; }
  bool operator!=(const mint &x) { return a != x.a; }
  mint operator+(const mint &x) {
    mint ret;
    ret.a = a + x.a;
    if (ret.a > mod)
      ret.a -= mod;
    return ret;
  }
  mint operator-(const mint &x) {
    mint ret;
    ret.a = a - x.a;
    if (ret.a < 0)
      ret.a += mod;
    return ret;
  }
  mint operator*(const mint &x) {
    mint ret;
    ret.a = a * x.a % mod;
    return ret;
  }
  mint operator/(mint &x) {
    mint ret;
    ret.a = a * x.inv().a % mod;
    return ret;
  }
  mint operator^(lint n) {
    mint ret = 1, x = a;
    while (n) {
      if (n & 1)
        (ret.a *= x.a) %= mod;
      (x.a *= x.a) %= mod;
      n >>= 1;
    }
    return ret;
  }
  mint &operator+=(const mint &x) {
    a += x.a;
    if (a >= mod)
      a -= mod;
    return *this;
  }
  mint &operator-=(const mint &x) {
    a -= x.a;
    if (a < 0)
      a += mod;
    return *this;
  }
  mint &operator*=(const mint &x) {
    (a *= x.a) %= mod;
    return *this;
  }
  mint &operator/=(mint &x) {
    (a *= x.inv().a) %= mod;
    return *this;
  }
  mint &operator^=(lint n) {
    lint ret = 1;
    while (n) {
      if (n & 1)
        (ret *= a) %= mod;
      (a *= a) %= mod;
      n >>= 1;
    }
    a = ret;
    return *this;
  }
};

vector<mint> fact;
vector<mint> revfact;

void setfact(int n) {
  fact.resize(n + 1);
  revfact.resize(n + 1);
  fact[0] = 1;
  rep(i, n) fact[i + 1] = fact[i] * mint(i + 1);

  revfact[n] = fact[n].inv();
  for (int i = n - 1; i >= 0; i--)
    revfact[i] = revfact[i + 1] * mint(i + 1);
}

mint getC(int n, int r) {
  if (n < r)
    return 0;
  return fact[n] * revfact[r] * revfact[n - r];
}

template <typename T> struct SegTree {
  using F = function<T(T, T)>;
  int sz = 1;
  T et;
  F f, g;
  vector<T> node;

  SegTree(int sz_, T et_, F f_, F g_) : et(et_), f(f_), g(g_) {
    while (sz < sz_)
      sz <<= 1;
    node.resize(sz << 1, et);
  }

  void update(int i, T x) {
    i += sz;
    node[i] = g(node[i], x);
    i >>= 1;
    while (i) {
      node[i] = f(node[i << 1], node[(i << 1) + 1]);
      i >>= 1;
    }
  }

  T query(int l, int r) {
    T vl = et, vr = et;
    for (l += sz, r += sz; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        vl = f(vl, node[l++]);
      if (r & 1)
        vr = f(node[--r], vr);
    }
    return f(vl, vr);
  }
};

int main() {
  int n, K;
  scanf("%d%d", &n, &K);
  int a[n];
  vector<int> v(n);
  rep(i, n) {
    scanf("%d", &a[i]);
    v[i] = a[i];
  }
  sort(v.begin(), v.end());
  v.erase(unique(v.begin(), v.end()), v.end());
  SegTree<int> st(v.size(), 0, plus<>(), plus<>());
  rep(i, n) {
    int idx = lower_bound(v.begin(), v.end(), a[i]) - v.begin();
    st.update(idx, 1);
  }
  setfact(n);
  mint ans = 0;
  sort(a, a + n);
  for (int t : v) {
    auto lb = lower_bound(a, a + n, t);
    auto ub = upper_bound(a, a + n, t);
    int x = lb - a;
    int y = ub - lb;
    int z = (a + n) - ub;
    auto P = (getC(y + z, K) - getC(z, K)) * mint((t + mod) % mod);
    auto Q = (getC(x + y, K) - getC(x, K)) * mint((t + mod) % mod);
    ans += Q - P;
    // printf("%lld %lld\n", P.a, Q.a);
    // printf("%d %d %d %d\n", x, y, z, t);
  }
  printf("%lld\n", ans.a);
}