Skip to content

3317. Find the Number of Possible Ways for an Event 👍

  • Time: $O(n \cdot \max(n, x))$
  • Space: $O(n \cdot \max(n, x))$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
class Solution {
 public:
  int numberOfWays(int n, int x, int y) {
    const int maxStages = min(n, x);
    const auto [fact, invFact] = getFactAndInvFact(max(n, x));
    const vector<vector<int>> stirling = getStirling(n, maxStages);
    int ans = 0;

    for (int k = 1; k <= maxStages; ++k) {
      // 1. Choose `k` stages from `x` stages.
      long events = nCk(x, k, fact, invFact);
      // 2. Partition `n` performers into `k` stages.
      events = events * stirling[n][k] % kMod;
      // 3. Permute `k` stages.
      events = events * fact[k] % kMod;
      // 4. Score `k` stages with score in the range [1, y], so y^k ways.
      events = events * modPow(y, k) % kMod;
      ans = (ans + events) % kMod;
    }

    return ans;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  pair<vector<long>, vector<long>> getFactAndInvFact(int n) {
    vector<long> fact(n + 1);
    vector<long> invFact(n + 1);
    vector<long> inv(n + 1);
    fact[0] = invFact[0] = 1;
    inv[0] = inv[1] = 1;
    for (int i = 1; i <= n; ++i) {
      if (i >= 2)
        inv[i] = kMod - kMod / i * inv[kMod % i] % kMod;
      fact[i] = fact[i - 1] * i % kMod;
      invFact[i] = invFact[i - 1] * inv[i] % kMod;
    }
    return {fact, invFact};
  }

  int nCk(int n, int k, const vector<long>& fact, const vector<long>& invFact) {
    return fact[n] * invFact[k] % kMod * invFact[n - k] % kMod;
  }

  // Returns a 2D array stirling, where stirling[i][j] := the number of ways to
  // partition a set of i objects into j non-empty subsets.
  //
  // https://en.wikipedia.org/wiki/Stirling_numbers_of_the_second_kind
  vector<vector<int>> getStirling(int n, int k) {
    vector<vector<int>> stirling(n + 1, vector<int>(k + 1));
    stirling[0][0] = 1;
    for (int i = 1; i <= n; ++i) {
      stirling[i][1] = 1;
      for (int j = 2; j <= min(i, k); ++j)
        stirling[i][j] = (static_cast<long>(j) * stirling[i - 1][j] +
                          stirling[i - 1][j - 1]) %
                         kMod;
    }
    return stirling;
  }

  long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x % kMod, (n - 1)) % kMod;
    return modPow(x * x % kMod, (n / 2)) % kMod;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class Solution {
  public int numberOfWays(int n, int x, int y) {
    final int maxStages = Math.min(n, x);
    long[][] factAndInvFact = getFactAndInvFact(Math.max(n, x));
    long[] fact = factAndInvFact[0];
    long[] invFact = factAndInvFact[1];
    int[][] stirling = getStirling(n, maxStages);
    int ans = 0;

    for (int k = 1; k <= maxStages; ++k) {
      // 1. Choose `k` stages from `x` stages.
      long events = nCk(x, k, fact, invFact);
      // 2. Partition `n` performers into `k` stages.
      events = events * stirling[n][k] % MOD;
      // 3. Permute `k` stages.
      events = events * fact[k] % MOD;
      // 4. Score `k` stages with score in the range [1, y], so y^k ways.
      events = events * modPow(y, k) % MOD;
      ans = (int) ((ans + events) % MOD);
    }

    return ans;
  }

  private static final int MOD = 1_000_000_007;

  private long[][] getFactAndInvFact(int n) {
    long[] fact = new long[n + 1];
    long[] invFact = new long[n + 1];
    long[] inv = new long[n + 1];
    fact[0] = invFact[0] = 1;
    inv[0] = inv[1] = 1;
    for (int i = 1; i <= n; ++i) {
      if (i >= 2)
        inv[i] = MOD - MOD / i * inv[MOD % i] % MOD;
      fact[i] = fact[i - 1] * i % MOD;
      invFact[i] = invFact[i - 1] * inv[i] % MOD;
    }
    return new long[][] {fact, invFact};
  }

  private int nCk(int n, int k, long[] fact, long[] invFact) {
    return (int) (fact[n] * invFact[k] % MOD * invFact[n - k] % MOD);
  }

  // Returns a 2D array stirling, where stirling[i][j] := the number of ways to
  // partition a set of i objects into j non-empty subsets.
  //
  // https://en.wikipedia.org/wiki/Stirling_numbers_of_the_second_kind
  private int[][] getStirling(int n, int k) {
    int[][] stirling = new int[n + 1][k + 1];
    stirling[0][0] = 1;
    for (int i = 1; i <= n; ++i) {
      stirling[i][1] = 1;
      for (int j = 2; j <= Math.min(i, k); ++j)
        stirling[i][j] = (int) (((long) j * stirling[i - 1][j] + stirling[i - 1][j - 1]) % MOD);
    }
    return stirling;
  }

  private long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x, n - 1) % MOD;
    return modPow(x * x % MOD, n / 2);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution:
  def numberOfWays(self, n: int, x: int, y: int) -> int:
    MOD = 1_000_000_007

    @functools.lru_cache(None)
    def fact(i: int) -> int:
      return 1 if i <= 1 else i * fact(i - 1) % MOD

    @functools.lru_cache(None)
    def inv(i: int) -> int:
      return pow(i, MOD - 2, MOD)

    @functools.lru_cache(None)
    def nCk(n: int, k: int) -> int:
      return fact(n) * inv(fact(k)) * inv(fact(n - k)) % MOD

    @functools.lru_cache(None)
    def stirling(n: int, k: int) -> int:
      """
      Returns the number of ways to partition a set of n objects into k
      non-empty subsets.

      https://en.wikipedia.org/wiki/Stirling_numbers_of_the_second_kind
      """
      if k == 0 or n < k:
        return 0
      if k == 1 or n == k:
        return 1
      return (k * stirling(n - 1, k) + stirling(n - 1, k - 1)) % MOD

    # 1. Choose `k` stages from `x` stages.
    # 2. Partition `n` performers into `k` stages.
    # 3. Permute `k` stages.
    # 4. Score `k` stages with score in the range [1, y], so y^k ways.
    return sum(nCk(x, k) * stirling(n, k) * fact(k) * pow(y, k, MOD) % MOD
               for k in range(1, min(n, x) + 1)) % MOD