Good Pairs  (Codeforces 1762F)
Tags: DP, binary search

You are given an array $$$a$$$ consisting of $$$n$$$ integers and an integer $$$k$$$.A pair $$$(l,r)$$$ is good if there exists a sequence of indices $$$i_1, i_2, \dots, i_m$$$ such that   $$$i_1=l$$$ and $$$i_m=r$$$;  $$$i_j < i_{j+1}$$$ for all $$$1 \leq j < m$$$; and  $$$|a_{i_j}-a_{i_{j+1}}| \leq k$$$ for all $$$1 \leq j < m$$$. Find the number of pairs $$$(l,r)$$$ ($$$1 \leq l \leq r \leq n$$$) that are good.

Input: Each test contains multiple test cases. The first line contains a single integer $$$t$$$ ($$$1 \leq t \leq 10^5$$$) — the number of test cases. The description of the test cases follows.The first line of each test case contains two space-separated integers $$$n$$$ and $$$k$$$ ($$$1 \leq n \leq 5 \cdot 10^5$$$; $$$0 \leq k \leq 10^5$$$) — the length of the array $$$a$$$ and the integer $$$k$$$.The second line of each test case contains $$$n$$$ space-separated integers $$$a_1,a_2,\ldots,a_n$$$ ($$$1 \leq a_i \leq 10^5$$$) — representing the array $$$a$$$.It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$5 \cdot 10^5$$$.

Output: For each test case, print the number of good pairs.

Note: In the first test case, good pairs are $$$(1,1)$$$, $$$(1,2)$$$, $$$(1,3)$$$, $$$(2,2)$$$, $$$(2,3)$$$, and $$$(3,3)$$$.In the second test case, good pairs are $$$(1,1)$$$, $$$(1,3)$$$, $$$(1,4)$$$, $$$(2,2)$$$, $$$(2,3)$$$, $$$(2,4)$$$, $$$(3,3)$$$, $$$(3,4)$$$ and $$$(4,4)$$$. Pair $$$(1,4)$$$ is good because there exists a sequence of indices $$$1, 3, 4$$$ which satisfy the given conditions.

Sample Inputs:
4
3 0
1 1 1
4 2
4 8 6 8
6 4
7 2 5 8 3 8
20 23
110 57 98 14 20 1 60 82 108 37 82 73 8 46 38 35 106 115 58 112

Sample Outputs:
6
9
18
92

Solution (Python):
import re
import functools
import random
import sys
import os
import math
from collections import Counter, defaultdict, deque
from functools import lru_cache, reduce
from itertools import accumulate, combinations, permutations
from heapq import nsmallest, nlargest, heappushpop, heapify, heappop, heappush
from io import BytesIO, IOBase
from copy import deepcopy
import threading
from typing import *
from operator import add, xor, mul, ior, iand, itemgetter
import bisect
BUFSIZE = 4096
inf = float('inf')

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')
sys.stdin = IOWrapper(sys.stdin)
sys.stdout = IOWrapper(sys.stdout)
input = lambda : sys.stdin.readline().rstrip('\r\n')

def I():
    return input()

def II():
    return int(input())

def MII():
    return map(int, input().split())

def LI():
    return list(input().split())

def LII():
    return list(map(int, input().split()))

def GMI():
    return map(lambda x: int(x) - 1, input().split())

def LGMI():
    return list(map(lambda x: int(x) - 1, input().split()))
from types import GeneratorType

def bootstrap(f, stack=[]):

    def wrappedfunc(*args, **kwargs):
        if stack:
            return f(*args, **kwargs)
        else:
            to = f(*args, **kwargs)
            while True:
                if type(to) is GeneratorType:
                    stack.append(to)
                    to = next(to)
                else:
                    stack.pop()
                    if not stack:
                        break
                    to = stack[-1].send(to)
            return to
    return wrappedfunc

class SegmentTree:

    def __init__(self, n, merge):
        self.n = n
        self.tree = [5 * 10 ** 5] * (2 * self.n)
        self._merge = merge

    def query(self, ql, qr):
        lans = rans = 5 * 10 ** 5
        ql += self.n
        qr += self.n + 1
        while ql < qr:
            if ql % 2:
                lans = self._merge(lans, self.tree[ql])
                ql += 1
            if qr % 2:
                qr -= 1
                rans = self._merge(rans, self.tree[qr])
            ql //= 2
            qr //= 2
        return self._merge(lans, rans)

    def update(self, index, value):
        index += self.n
        self.tree[index] = value
        while index:
            index //= 2
            self.tree[index] = self._merge(self.tree[2 * index], self.tree[2 * index + 1])

class FenwickTree:

    def __init__(self, n):
        self.n = n
        self.bit = [0] * n

    def sum(self, r):
        res = 0
        while r >= 0:
            res += self.bit[r]
            r = (r & r + 1) - 1
        return res

    def rsum(self, l, r):
        return self.sum(r) - self.sum(l - 1)

    def add(self, idx, delta):
        while idx < self.n:
            self.bit[idx] += delta
            idx = idx | idx + 1

def solve(nums, k):
    n = len(nums)
    dp = [0] * n
    for i in range(n - 1, -1, -1):
        if nums[i] < 10 ** 5:
            idx = seg.query(nums[i] + 1, min(10 ** 5 - 1, nums[i] + k))
            if idx < n:
                dp[i] = dp[idx] + fenwick.rsum(nums[i] + 1, nums[idx])
        fenwick.add(nums[i], 1)
        seg.update(nums[i], i)
    ans = sum(dp)
    for i in range(n):
        fenwick.add(nums[i], -1)
        seg.update(nums[i], 5 * 10 ** 5)
    return ans
t = II()
seg = SegmentTree(10 ** 5, lambda x, y: x if x <= y else y)
fenwick = FenwickTree(10 ** 5)
count = [0] * 10 ** 5
for _ in range(t):
    (n, k) = MII()
    nums = LGMI()
    ans = 0
    for num in nums:
        count[num] += 1
        ans += count[num]
    if k:
        ans += solve(nums, k) + solve(nums[::-1], k)
    for num in nums:
        count[num] -= 1
    print(ans)

Solution (Python):
import re
import functools
import random
import sys
import os
import math
from collections import Counter, defaultdict, deque
from functools import lru_cache, reduce
from itertools import accumulate, combinations, permutations
from heapq import nsmallest, nlargest, heappushpop, heapify, heappop, heappush
from io import BytesIO, IOBase
from copy import deepcopy
import threading
from typing import *
from operator import add, xor, mul, ior, iand, itemgetter
import bisect
BUFSIZE = 4096
inf = float('inf')

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')
sys.stdin = IOWrapper(sys.stdin)
sys.stdout = IOWrapper(sys.stdout)
input = lambda : sys.stdin.readline().rstrip('\r\n')

def I():
    return input()

def II():
    return int(input())

def MII():
    return map(int, input().split())

def LI():
    return list(input().split())

def LII():
    return list(map(int, input().split()))

def GMI():
    return map(lambda x: int(x) - 1, input().split())

def LGMI():
    return list(map(lambda x: int(x) - 1, input().split()))
from types import GeneratorType

def bootstrap(f, stack=[]):

    def wrappedfunc(*args, **kwargs):
        if stack:
            return f(*args, **kwargs)
        else:
            to = f(*args, **kwargs)
            while True:
                if type(to) is GeneratorType:
                    stack.append(to)
                    to = next(to)
                else:
                    stack.pop()
                    if not stack:
                        break
                    to = stack[-1].send(to)
            return to
    return wrappedfunc

class SegmentTree:

    def __init__(self, n, merge):
        self.n = n
        self.tree = [5 * 10 ** 5] * (2 * self.n)
        self._merge = merge

    def query(self, ql, qr):
        lans = rans = 5 * 10 ** 5
        ql += self.n
        qr += self.n + 1
        while ql < qr:
            if ql % 2:
                lans = self._merge(lans, self.tree[ql])
                ql += 1
            if qr % 2:
                qr -= 1
                rans = self._merge(rans, self.tree[qr])
            ql //= 2
            qr //= 2
        return self._merge(lans, rans)

    def update(self, index, value):
        index += self.n
        self.tree[index] = value
        while index:
            index //= 2
            self.tree[index] = self._merge(self.tree[2 * index], self.tree[2 * index + 1])

class FenwickTree:

    def __init__(self, n):
        self.n = n
        self.bit = [0] * n

    def sum(self, r):
        res = 0
        while r >= 0:
            res += self.bit[r]
            r = (r & r + 1) - 1
        return res

    def rsum(self, l, r):
        return self.sum(r) - self.sum(l - 1)

    def add(self, idx, delta):
        while idx < self.n:
            self.bit[idx] += delta
            idx = idx | idx + 1

def solve(nums, k):
    n = len(nums)
    dp = [0] * n
    for i in range(n - 1, -1, -1):
        if nums[i] < 10 ** 5:
            idx = seg.query(nums[i] + 1, min(10 ** 5 - 1, nums[i] + k))
            if idx < n:
                dp[i] = dp[idx] + fenwick.rsum(nums[i] + 1, nums[idx])
        fenwick.add(nums[i], 1)
        seg.update(nums[i], i)
    ans = sum(dp)
    for i in range(n):
        fenwick.add(nums[i], -1)
        seg.update(nums[i], 5 * 10 ** 5)
    return ans
t = II()
seg = SegmentTree(10 ** 5, lambda x, y: x if x <= y else y)
fenwick = FenwickTree(10 ** 5)
count = [0] * 10 ** 5
for _ in range(t):
    (n, k) = MII()
    nums = LGMI()
    ans = 0
    for num in nums:
        count[num] += 1
        ans += count[num]
    if k:
        ans += solve(nums, k) + solve(nums[::-1], k)
    for num in nums:
        count[num] -= 1
    print(ans)

Solution (Python):
import re
import functools
import random
import sys
import os
import math
from collections import Counter, defaultdict, deque
from functools import lru_cache, reduce
from itertools import accumulate, combinations, permutations
from heapq import nsmallest, nlargest, heappushpop, heapify, heappop, heappush
from io import BytesIO, IOBase
from copy import deepcopy
import threading
from typing import *
from operator import add, xor, mul, ior, iand, itemgetter
import bisect
BUFSIZE = 4096
inf = float('inf')

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')
sys.stdin = IOWrapper(sys.stdin)
sys.stdout = IOWrapper(sys.stdout)
input = lambda : sys.stdin.readline().rstrip('\r\n')

def I():
    return input()

def II():
    return int(input())

def MII():
    return map(int, input().split())

def LI():
    return list(input().split())

def LII():
    return list(map(int, input().split()))

def GMI():
    return map(lambda x: int(x) - 1, input().split())

def LGMI():
    return list(map(lambda x: int(x) - 1, input().split()))
from types import GeneratorType

def bootstrap(f, stack=[]):

    def wrappedfunc(*args, **kwargs):
        if stack:
            return f(*args, **kwargs)
        else:
            to = f(*args, **kwargs)
            while True:
                if type(to) is GeneratorType:
                    stack.append(to)
                    to = next(to)
                else:
                    stack.pop()
                    if not stack:
                        break
                    to = stack[-1].send(to)
            return to
    return wrappedfunc

class SegmentTree:

    def __init__(self, n, merge):
        self.n = n
        self.tree = [5 * 10 ** 5] * (2 * self.n)
        self._merge = merge

    def query(self, ql, qr):
        lans = rans = 5 * 10 ** 5
        ql += self.n
        qr += self.n + 1
        while ql < qr:
            if ql % 2:
                lans = self._merge(lans, self.tree[ql])
                ql += 1
            if qr % 2:
                qr -= 1
                rans = self._merge(rans, self.tree[qr])
            ql //= 2
            qr //= 2
        return self._merge(lans, rans)

    def update(self, index, value):
        index += self.n
        self.tree[index] = value
        while index:
            index //= 2
            self.tree[index] = self._merge(self.tree[2 * index], self.tree[2 * index + 1])

class FenwickTree:

    def __init__(self, n):
        self.n = n
        self.bit = [0] * n

    def sum(self, r):
        res = 0
        while r >= 0:
            res += self.bit[r]
            r = (r & r + 1) - 1
        return res

    def rsum(self, l, r):
        return self.sum(r) - self.sum(l - 1)

    def add(self, idx, delta):
        while idx < self.n:
            self.bit[idx] += delta
            idx = idx | idx + 1

def solve(nums, k):
    n = len(nums)
    dp = [0] * n
    for i in range(n - 1, -1, -1):
        if nums[i] < 10 ** 5:
            idx = seg.query(nums[i] + 1, min(10 ** 5 - 1, nums[i] + k))
            if idx < n:
                dp[i] = dp[idx] + fenwick.rsum(nums[i] + 1, nums[idx])
        fenwick.add(nums[i], 1)
        seg.update(nums[i], i)
    ans = sum(dp)
    for i in range(n):
        fenwick.add(nums[i], -1)
        seg.update(nums[i], 5 * 10 ** 5)
    return ans
t = II()
seg = SegmentTree(10 ** 5, lambda x, y: x if x <= y else y)
fenwick = FenwickTree(10 ** 5)
count = [0] * 10 ** 5
for _ in range(t):
    (n, k) = MII()
    nums = LGMI()
    ans = 0
    for num in nums:
        count[num] += 1
        ans += count[num]
    if k:
        ans += solve(nums, k) + solve(nums[::-1], k)
    for num in nums:
        count[num] -= 1
    print(ans)