codeforces-go/copypasta/math_fwt.go at master · EndlessCheng/codeforces-go · GitHub

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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
package copypasta

import "slices"

/* 快速沃尔什变换 fast Walsh–Hadamard transform, FWT, FWHT

用于解决对下标执行【位运算卷积】的问题
一个常见的应用场景是对频率数组求 FWT
例如把 a 里面所有元素的出现次数统计到一个 cnt 数组中，cnt[x] 表示 x 在 a 中的出现次数
然后求 FWT(cnt, cnt)，就得到了 a 中任意两个元素的【异或/或/和】的结果的出现次数（特别地，这得到了任意两个数的【异或/或/和】的最大值）
https://leetcode.cn/problems/number-of-unique-xor-triplets-ii/

又例如，求一个数组的三个元素的最大异或和，在值域不大的情况下，
可以先求出该数组的频率数组与频率数组的 FWT，即得到两个元素的所有异或和（及组成该异或和的元素对数），
然后枚举两元素异或和，在原数组的异或字典树上查询最大异或和
具体到名称，OR 上的 FWT 也叫 fast zeta transform，AND 上的 FWT 也叫 fast mobius transform
https://en.wikipedia.org/wiki/Fast_Walsh%E2%80%93Hadamard_transform
https://oi-wiki.org/math/poly/fwt/
OR Convolution for Common People https://codeforces.com/blog/entry/115438
https://www.luogu.com.cn/blog/CJL/fwt-xue-xi-bi-ji
https://www.luogu.com.cn/blog/command-block/wei-yun-suan-juan-ji-yu-ji-kuo-zhan
一些习题 https://blog.csdn.net/weixin_38686780/article/details/81912853

模板题 https://www.luogu.com.cn/problem/P4717
LC1803 https://leetcode.cn/problems/count-pairs-with-xor-in-a-range/
与仙人掌结合 https://codeforces.com/problemset/problem/1218/D
https://codeforces.com/problemset/problem/662/C 2600
 - https://atcoder.jp/contests/abc396/tasks/abc396_g
https://codeforces.com/problemset/problem/914/G 2600
https://atcoder.jp/contests/abc212/tasks/abc212_h

EXTRA: 子集卷积
todo https://www.cnblogs.com/yijan/p/12387352.html
 https://www.luogu.com.cn/problem/P6097
*/

// 取模的写法见 https://www.luogu.com.cn/record/51397587
func fwtOR(a []int, op int) {
	n := len(a)
	for l, k := 2, 1; l <= n; l, k = l<<1, k<<1 {
		for i := 0; i < n; i += l {
			for j := 0; j < k; j++ {
				a[i+j+k] += a[i+j] * op // 注意负数
			}
		}
	}
}

func fwtAND(a []int, op int) {
	n := len(a)
	for l, k := 2, 1; l <= n; l, k = l<<1, k<<1 {
		for i := 0; i < n; i += l {
			for j := 0; j < k; j++ {
				a[i+j] += a[i+j+k] * op // 注意负数
			}
		}
	}
}

func fwtXOR(a []int, op int) {
	n := len(a)
	for l, k := 2, 1; l <= n; l, k = l<<1, k<<1 {
		for i := 0; i < n; i += l {
			for j := 0; j < k; j++ {
				// 若题目没有取模，IFWT 时 *op 改成 >>1
				a[i+j], a[i+j+k] = (a[i+j]+a[i+j+k])*op, (a[i+j]-a[i+j+k])*op // 注意负数
			}
		}
	}
}

// 求 OR 和 AND 时 invOp = -1
// 求 XOR 时 invOp = inv(2)，如果没有取模，就是 >>0 和 >>1
func fwt(a, b []int, fwtFunc func([]int, int), invOp int) []int {
	// 不修改原始数组
	a = slices.Clone(a)
	b = slices.Clone(b)
	fwtFunc(a, 1)
	fwtFunc(b, 1)
	for i, v := range b {
		a[i] *= v // mod
	}
	fwtFunc(a, invOp)
	return a
}

// 注：若代码性能瓶颈在 FWT 上，可以通过以下方式消除比较慢的乘法和取模
// （CF 旧版本用的 386）
// 优化前 1575ms https://codeforces.com/contest/1218/submission/118700754
// 优化后  748ms https://codeforces.com/contest/1218/submission/118704484
// （CF 更新成 amd64 之后）
// 优化前 779ms https://codeforces.com/contest/1218/submission/225480738
// 优化后 545ms https://codeforces.com/contest/1218/submission/225480844
func add(a, b int) int {
	a += b
	if a >= mod {
		a -= mod
	}
	return a
}

func sub_(a, b int) int {
	a -= b
	if a < 0 {
		a += mod
	}
	return a
}

func div2(a int) int {
	if a&1 > 0 {
		a += mod
	}
	return a >> 1
}

// 将 (-mod,mod) 范围内的 a 变成 [0,mod) 范围内
// 原理是负数右移会不断补 1，所以最后二进制都是 1，因此返回值等价于 a+_mod
// 而对于非负数，右移后二进制全为 0，所以返回结果仍然是 a
func norm32(a int32) int32 {
	return a + a>>31&mod
}

func norm64(a int64) int64 {
	return a + a>>63&mod
}