Skip to content

3098. Find the Sum of Subsequence Powers 👍

Approach 1: Memorize by indices that can lead to the minimum difference

  • Time: $O(n^4k)$
  • Space: $O(n^3k)$
 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
class Solution {
 public:
  int sumOfPowers(vector<int>& nums, int k) {
    const int n = nums.size();
    ranges::sort(nums);
    vector<vector<vector<vector<int>>>> mem(
        n + 1, vector<vector<vector<int>>>(
                   n + 1, vector<vector<int>>(n + 1, vector<int>(k + 1, -1))));
    return sumOfPowers(nums, 0, k, -1, -1, -1, mem);
  }

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

  // Returns the sum of powers of all subsequences of nums[i..n) which
  // have length equal to k, where `lastPickedIndex` is the index of the last
  // picked number and nums[secondIndex] - nums[firstIndex] is the minimum power
  // so far.
  int sumOfPowers(const vector<int>& nums, int i, int k, int lastPickedIndex,
                  int firstIndex, int secondIndex,
                  vector<vector<vector<vector<int>>>>& mem) {
    if (k == 0)
      return nums[secondIndex] - nums[firstIndex];
    if (i == nums.size())
      return 0;
    const int a = hash(lastPickedIndex);
    const int b = hash(firstIndex);
    const int c = hash(secondIndex);
    if (mem[a][b][c][k] != -1)
      return mem[a][b][c][k];
    int newFirstIndex = firstIndex;
    int newSecondIndex = secondIndex;
    if (firstIndex == -1) {
      newFirstIndex = i;
    } else if (secondIndex == -1) {
      newSecondIndex = i;
    } else if (nums[i] - nums[lastPickedIndex] <
               nums[secondIndex] - nums[firstIndex]) {
      newFirstIndex = lastPickedIndex;
      newSecondIndex = i;
    }
    const int pick =
        sumOfPowers(nums, i + 1, k - 1, i, newFirstIndex, newSecondIndex, mem);
    const int skip = sumOfPowers(nums, i + 1, k, lastPickedIndex, firstIndex,
                                 secondIndex, mem);
    return mem[a][b][c][k] = (pick + skip) % kMod;
  }

  constexpr int hash(int x) {
    return x + 1;
  }
};

 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
class Solution {
  public int sumOfPowers(int[] nums, int k) {
    final int n = nums.length;
    Arrays.sort(nums);
    Integer[][][][] mem = new Integer[n + 1][n + 1][n + 1][k + 1];
    return sumOfPowers(nums, 0, k, -1, -1, -1, mem);
  }

  private static final int kMod = 1_000_000_007;

  // Returns the sum of powers of all subsequences of nums[i..n) which
  // have length equal to k, where `lastPickedIndex` is the index of the last
  // picked number and nums[secondIndex] - nums[firstIndex] is the minimum power
  // so far.
  private int sumOfPowers(int[] nums, int i, int k, int lastPickedIndex, int firstIndex,
                          int secondIndex, Integer[][][][] mem) {
    if (k == 0)
      return nums[secondIndex] - nums[firstIndex];
    if (i == nums.length)
      return 0;
    final int a = hash(lastPickedIndex);
    final int b = hash(firstIndex);
    final int c = hash(secondIndex);
    if (mem[a][b][c][k] != null)
      return mem[a][b][c][k];
    int newFirstIndex = firstIndex;
    int newSecondIndex = secondIndex;
    if (firstIndex == -1) {
      newFirstIndex = i;
    } else if (secondIndex == -1) {
      newSecondIndex = i;
    } else if (nums[i] - nums[lastPickedIndex] < nums[secondIndex] - nums[firstIndex]) {
      newFirstIndex = lastPickedIndex;
      newSecondIndex = i;
    }
    final int pick = sumOfPowers(nums, i + 1, k - 1, i, newFirstIndex, newSecondIndex, mem);
    final int skip = sumOfPowers(nums, i + 1, k, lastPickedIndex, firstIndex, secondIndex, mem);
    return mem[a][b][c][k] = (pick + skip) % kMod;
  }

  private int hash(int x) {
    return x + 1;
  }
}

 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
class Solution:
  def sumOfPowers(self, nums: list[int], k: int) -> int:
    kMod = 1_000_000_007
    nums.sort()

    @functools.lru_cache(None)
    def dp(
            i: int, k: int, lastPickedIndex: int, firstIndex: int,
            secondIndex: int) -> int:
      if k == 0:
        return nums[secondIndex] - nums[firstIndex]
      if i == len(nums):
        return 0
      newFirstIndex = firstIndex
      newSecondIndex = secondIndex
      if firstIndex == -1:
        newFirstIndex = i
      elif secondIndex == -1:
        newSecondIndex = i
      elif nums[i] - nums[lastPickedIndex] < nums[secondIndex] - nums[firstIndex]:
        newFirstIndex = lastPickedIndex
        newSecondIndex = i
      pick = dp(i + 1, k - 1, i, newFirstIndex, newSecondIndex)
      skip = dp(i + 1, k, lastPickedIndex, firstIndex, secondIndex)
      return (pick + skip) % kMod

    return dp(0, k, -1, -1, -1)

Approach 2: Memorize by the minimum difference

  • Time: $O(n^4k)$
  • Space: $O(n^3k)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution:
  def sumOfPowers(self, nums: list[int], k: int) -> int:
    kMod = 1_000_000_007
    n = len(nums)

    nums.sort()

    @functools.lru_cache(None)
    def dp(i: int, k: int, lastPickIndex: int, minDiff: int) -> int:
      if k == 0:
        return minDiff
      if i == n:
        return 0
      newMinDiff = (minDiff if lastPickIndex == - 1
                    else min(minDiff, nums[i] - nums[lastPickIndex]))
      pick = dp(i + 1, k - 1, i, newMinDiff)
      skip = dp(i + 1, k, lastPickIndex, minDiff)
      return (pick + skip) % kMod

    return dp(0, k, -1, math.inf)