Maximum Sum of Alternating Subsequence With Distance at Least K solution and explanation

class FenwickTree:
    def __init__(self, n):
        self.n = n
        self.tree = [0] * (n + 1)

    def update(self, index: int, val: int) -> None:
        while index <= self.n:
            self.tree[index] = max(self.tree[index], val)
            index += index & (-index)  # 往后更新

    def preSum(self, pos):
        # 按照预期的方式求前缀最大值
        ans = 0
        while pos >= 1:
            ans = max(ans, self.tree[pos])
            pos -= pos & (-pos)
        return ans


class Solution:
    def maxAlternatingSum(self, nums: list[int], k: int) -> int:
        stl = sorted(set(nums))  # 将nums中不同的数字进行排序
        rank = {
            v: i + 1 for i, v in enumerate(stl)
        }  # 将nums中的值快速转换成stl中的索引
        fwt0 = FenwickTree(len(stl))
        fwt1 = FenwickTree(len(stl))

        n = len(nums)
        dp = [[0, 0] for _ in range(n)]
        res = nums[0]
        for i in range(n):
            dp[i][0] = dp[i][1] = nums[i]
            if i >= k:
                indx = rank[nums[i]]  # 找到nums[i]在stl中的索引
                dp[i][1] = max(
                    dp[i][1], fwt0.preSum(indx - 1) + nums[i]
                )  # indx-1即表示小于nums[i]的部分
                dp[i][0] = max(
                    dp[i][0], fwt1.preSum(len(stl) - indx) + nums[i]
                )  # len(stl)-indx即表示在倒序列表中大于nums[i]的部分

            if i - k + 1 >= 0:
                indx = rank[nums[i - k + 1]]
                fwt0.update(indx, dp[i - k + 1][0])  # 在正序列表中更新i-k+1位置的值
                fwt1.update(
                    len(stl) - indx + 1, dp[i - k + 1][1]
                )  # 在倒序列表中更新i-k+1位置的值

            res = max(res, dp[i][0], dp[i][1])  # 更新答案

        return res

class Solution {
    public long maxAlternatingSum(int[] nums, int k) {
        long maxSum = 0;
        int n = nums.length;
        int m = Arrays.stream(nums).max().getAsInt();
        long[][] dp = new long[n][2];
        SegmentTree[] sts = new SegmentTree[2];
        for (int j = 0; j < 2; j++) {
            sts[j] = new SegmentTree(m + 1);
        }
        for (int i = 0; i < n; i++) {
            if (i >= k) {
                sts[0].update(nums[i - k], dp[i - k][0]);
                sts[1].update(nums[i - k], dp[i - k][1]);
            }
            dp[i][0] = sts[1].getMax(0, nums[i] - 1) + nums[i];
            dp[i][1] = sts[0].getMax(nums[i] + 1, m) + nums[i];
            maxSum = Math.max(maxSum, Math.max(dp[i][0], dp[i][1]));
        }
        return maxSum;
    }
}

class SegmentTree {
    private int n;
    private long[] tree;

    public SegmentTree(int n) {
        this.n = n;
        this.tree = new long[n * 4];
    }

    public long getMax(int start, int end) {
        return getMax(start, end, 0, 0, n - 1);
    }

    public void update(int index, long value) {
        update(index, value, 0, 0, n - 1);
    }

    private long getMax(int rangeStart, int rangeEnd, int treeIndex, int treeStart, int treeEnd) {
        if (rangeStart > rangeEnd) {
            return 0;
        }
        if (rangeStart == treeStart && rangeEnd == treeEnd) {
            return tree[treeIndex];
        }
        int mid = treeStart + (treeEnd - treeStart) / 2;
        if (rangeEnd <= mid) {
            return getMax(rangeStart, rangeEnd, treeIndex * 2 + 1, treeStart, mid);
        } else if (rangeStart > mid) {
            return getMax(rangeStart, rangeEnd, treeIndex * 2 + 2, mid + 1, treeEnd);
        } else {
            return Math.max(getMax(rangeStart, mid, treeIndex * 2 + 1, treeStart, mid), getMax(mid + 1, rangeEnd, treeIndex * 2 + 2, mid + 1, treeEnd));
        }
    }

    private void update(int rangeIndex, long value, int treeIndex, int start, int end) {
        if (start == end) {
            tree[treeIndex] = value;
            return;
        }
        int mid = start + (end - start) / 2;
        if (rangeIndex <= mid) {
            update(rangeIndex, value, treeIndex * 2 + 1, start, mid);
        } else {
            update(rangeIndex, value, treeIndex * 2 + 2, mid + 1, end);
        }
        tree[treeIndex] = Math.max(tree[treeIndex * 2 + 1], tree[treeIndex * 2 + 2]);
    }
}

class Solution {
public:
    long long maxAlternatingSum(vector<int>& nums, int K) {
        int n = nums.size();

        int idx[n];
        map<int, int> mp;
        for (int x : nums) mp[x] = 1;
        int m = 0;
        for (auto &p : mp) p.second = ++m;
        for (int i = 0; i < n; i++) idx[i] = mp[nums[i]];

        const long long INF = 1e18;
        long long tree[2][m + 1];
        for (int i = 0; i < 2; i++) for (int j = 0; j <= m; j++) tree[i][j] = -INF;


        auto lb = [&](int x) { return x & (-x); };

        auto update = [&](long long *tree, int pos, long long val) {
            for (; pos <= m; pos += lb(pos)) tree[pos] = max(tree[pos], val);
        };

        auto query = [&](long long *tree, int pos) {
            long long ret = -INF;
            for (; pos; pos -= lb(pos)) ret = max(ret, tree[pos]);
            return ret;
        };


        long long ans = 0;
        long long f[n + 1][2];
        for (int i = 0; i <= n; i++) for (int j = 0; j < 2; j++) f[i][j] = -INF;
        for (int i = 1, j = 1; i <= n; i++) {
            while (i - j >= K) {
                update(tree[0], idx[j - 1], f[j][0]);
                update(tree[1], m + 1 - idx[j - 1], f[j][1]);
                j++;
            }
            f[i][0] = max(0LL, query(tree[1], m - idx[i - 1])) + nums[i - 1];
            f[i][1] = max(0LL, query(tree[0], idx[i - 1] - 1)) + nums[i - 1];
            ans = max({ans, f[i][0], f[i][1]});
        }
        return ans;
    }
};

type fenwick []int64

func (f fenwick) update(i int, val int64) {
	for ; i < len(f); i += i & -i {
		f[i] = max(f[i], val)
	}
}

func (f fenwick) preMax(i int) (res int64) {
	for ; i > 0; i &= i - 1 {
		res = max(res, f[i])
	}
	return
}

func maxAlternatingSum(nums []int, k int) (ans int64) {
	sorted := slices.Clone(nums)
	slices.Sort(sorted)
	sorted = slices.Compact(sorted)

	n := len(nums)
	fInc := make([]int64, n)
	fDec := make([]int64, n)

	m := len(sorted)
	inc := make(fenwick, m+1)
	dec := make(fenwick, m+1)

	for i, x := range nums {
		if i >= k {
			j := nums[i-k]
			inc.update(m-j, fInc[i-k])
			dec.update(j+1, fDec[i-k])
		}

		j := sort.SearchInts(sorted, x)
		nums[i] = j

		fInc[i] = dec.preMax(j) + int64(x)
		fDec[i] = inc.preMax(m-1-j) + int64(x)
		ans = max(ans, fInc[i], fDec[i])
	}

	return
}