2019-1B Problem C

Final solution

// package, import, etc...

func min(a, b int) int {
	if a < b {
		return a
	} else {
		return b
	}
}

func max(a, b int) int {
	if a > b {
		return a
	} else {
		return b
	}
}

func ConstructMaxTable(arr []int) [][]int {
	logArrLen := 0
	for arrLen := len(arr); arrLen > 0; arrLen >>= 1 {
		logArrLen++
	}
	maxTable := make([][]int, 0, logArrLen+1) // store indices, not values
	for currentSegLen := 1; currentSegLen <= len(arr); currentSegLen *= 2 {
		thisMaxRow := make([]int, len(arr)-currentSegLen+1)
		if currentSegLen == 1 {
			for i := 0; i < len(arr); i++ {
				thisMaxRow[i] = i
			}
		} else {
			lastMaxRow := maxTable[len(maxTable)-1]
			for i := 0; i < len(arr)-currentSegLen+1; i++ {
				indexA := lastMaxRow[i]
				indexB := lastMaxRow[i+currentSegLen/2]
				if arr[indexA] >= arr[indexB] { // perfer low index
					thisMaxRow[i] = indexA
				} else {
					thisMaxRow[i] = indexB
				}
			}
		}
		maxTable = append(maxTable, thisMaxRow)
	}
	return maxTable
}

// find the maximum element in range [start, end) and return its index, perfering lower index when equal.
func RangeMax(arr []int, maxTable [][]int, start, end int) int {
	rangeLen := uint(end - start)
	if rangeLen == 0 {
		return start
	}
	var rangeLevel uint = 0
	var rpower uint
	for rpower = 1; rpower < rangeLen; rpower <<= 1 {
		rangeLevel++
	}
	if (1 << rangeLevel) == rangeLen {
		return maxTable[rangeLevel][start]
	} else {
		rangeLevel--
		indexA := maxTable[rangeLevel][start]
		indexB := maxTable[rangeLevel][end-(1<<rangeLevel)]
		if arr[indexA] >= arr[indexB] {
			return indexA
		} else {
			return indexB
		}
	}
}

func start() {
  // read T, repeat test() T times...
}

var nTypes, K int

func test() {
	mustReadLineOfInts(&nTypes, &K)
	cSkills := mustReadLineOfIntsIntoArray()
	assert(len(cSkills) == nTypes)
	dSkills := mustReadLineOfIntsIntoArray()
	cMaxTable := ConstructMaxTable(cSkills)
	dMaxTable := ConstructMaxTable(dSkills)
	assert(len(dSkills) == nTypes)
	count := 0
	for i := 0; i < nTypes; i++ {
		currentChosenSkill := cSkills[i]
		// ! ! ! . . . ...
		stillChooseILeft := 0 // least index that still result in choosing i
		if i > 0 {
			stillChooseILeft = sort.Search(i, func(left int) bool {
				// still choose i?
				newMaxI := RangeMax(cSkills, cMaxTable, left, i)
				// Problem here:
				// return newMaxI <= currentChosenSkill
				return cSkills[newMaxI] <= currentChosenSkill
			})
		}
		goodEnoughLeft := sort.Search(i+1, func(left int) bool {
			dChoose := RangeMax(dSkills, dMaxTable, left, i+1)
			return dSkills[dChoose]-currentChosenSkill <= K
		})
		tooGoodLeft := sort.Search(i+1, func(left int) bool {
			dChoose := RangeMax(dSkills, dMaxTable, left, i+1)
			return !(currentChosenSkill-dSkills[dChoose] <= K)
		})
		if goodEnoughLeft == i+1 {
			continue
		}
		assert(goodEnoughLeft <= tooGoodLeft)
		stillChooseIRight := nTypes // first index that will result in not choosing i
		if i < nTypes-1 {
			stillChooseIRight = sort.Search(nTypes, func(right int) bool {
				if right <= i {
					return false
				}
				// Problem here:
				// newMaxI := RangeMax(cSkills, cMaxTable, i+1, right)
				newMaxI := RangeMax(cSkills, cMaxTable, i+1, right+1)
				return !(cSkills[newMaxI] < currentChosenSkill)
			})
		}
		firstNotGoodEnoughRight := sort.Search(nTypes, func(right int) bool {
			if right < i {
				return false
			}
			dChoose := RangeMax(dSkills, dMaxTable, i, right+1)
			return !(dSkills[dChoose]-currentChosenSkill <= K)
		})
		if firstNotGoodEnoughRight == i {
			continue
		}
		firstNotTooGoodRight := sort.Search(nTypes, func(right int) bool {
			if right < i {
				return false
			}
			dChoose := RangeMax(dSkills, dMaxTable, i, right+1)
			return currentChosenSkill-dSkills[dChoose] <= K
		})
		assert(firstNotTooGoodRight <= firstNotGoodEnoughRight)
		ans := (i - max(stillChooseILeft, goodEnoughLeft) + 1) * (min(firstNotGoodEnoughRight, stillChooseIRight) - i)
		if tooGoodLeft < firstNotTooGoodRight {
			ans -= (i - max(stillChooseILeft, tooGoodLeft) + 1) * (min(firstNotTooGoodRight, stillChooseIRight) - i)
		}
		assert(ans >= 0)
		count += ans
	}
	fmt.Fprintf(stdout, "%d\n", count)
}

// boilerplate omitted...

Solution that only works for the small input

// package, import, etc...

func start() {
  // read T, repeat test() T times...
}

func max(a, b int) int {
	if a > b {
		return a
	} else {
		return b
	}
}

func abs(x int) int {
	if x < 0 {
		return -x
	} else {
		return x
	}
}

func test() {
	var nTypes, K int
	mustReadLineOfInts(&nTypes, &K)
	cSkills := mustReadLineOfIntsIntoArray()
	assert(len(cSkills) == nTypes)
	dSkills := mustReadLineOfIntsIntoArray()
	assert(len(dSkills) == nTypes)
	count := 0
	for i := 0; i < nTypes; i++ {
		currentCMax := cSkills[i]
		currentDMax := dSkills[i]
		for j := i + 1; j <= nTypes; j++ {
			currentCMax = max(currentCMax, cSkills[j-1])
			currentDMax = max(currentDMax, dSkills[j-1])
			if abs(currentCMax-currentDMax) <= K {
				count++
			}
		}
	}
	fmt.Fprintf(stdout, "%d\n", count)
}

// boilerplate omitted...

Sample test cases given in question

6
4 0
1 1 1 8
8 8 8 8
3 0
0 1 1
1 1 0
1 0
3
3
5 0
0 8 0 8 0
4 0 4 0 4
3 0
1 0 0
0 1 2
5 2
1 2 3 4 5
5 5 5 5 10

Other files

1
6 1
8 13 12 7 7 13
0 0 1 3 2 11

Case #1: 0

// package, import, etc...

func main() {
	rawIn, err := os.OpenFile("generated.in", os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		panic(err)
	}
	in := bufio.NewWriter(rawIn)
	defer rawIn.Close()
	defer in.Flush()
	const t int = 100
	fmt.Fprintf(in, "%d\n", t)
	for i := 0; i < t; i++ {
		var N, K int
		N = rand.Intn(10) + 1
		K = rand.Intn(5)
		fmt.Fprintf(in, "%d %d\n", N, K)
		for i := 0; i < 2; i++ {
			for i := 0; i < N; i++ {
				if i != 0 {
					in.WriteByte(' ')
				}
				in.WriteString(strconv.Itoa(rand.Intn(15)))
			}
			in.WriteByte('\n')
		}
	}
}