【PTA数据结构 | C语言版】树状数组的操作

本专栏持续输出数据结构题目集，欢迎订阅。

题目

请编写程序，实现树状数组区间求前缀和、单点修改的操作。

输入格式：
输入首先给出一个正整数 n（2≤n<10^3），随后一行给出 n 个绝对值不超过 10^5 的整数。

输出格式：
第一行按存储顺序输出树状数组中元素；第二行按存储顺序输出前缀和数组中的元素。
同行数字间以 1 个空格分隔。为输出简单起见，每个数字后有一个空格。

输入样例：
15
15 14 13 12 11 10 9 8 7 6 5 4 3 2 1

输出样例：
15 29 13 54 11 21 9 92 7 13 5 22 3 5 1
15 29 42 54 65 75 84 92 99 105 110 114 117 119 120

代码

#include <stdio.h>
#include <stdlib.h>

// 树状数组结构
typedef struct {
    int *tree;  // 树状数组
    int *arr;   // 原始数组
    int size;   // 数组大小
} FenwickTree;

// 初始化树状数组
FenwickTree* initFenwickTree(int *data, int n) {
    FenwickTree *ft = (FenwickTree*)malloc(sizeof(FenwickTree));
    ft->size = n;
    ft->arr = (int*)malloc((n + 1) * sizeof(int));  // 1-based索引
    ft->tree = (int*)malloc((n + 1) * sizeof(int)); // 1-based索引
    
    // 复制原始数据（转换为1-based）
    for (int i = 1; i <= n; i++) {
        ft->arr[i] = data[i - 1];
        ft->tree[i] = 0;
    }
    
    // 构建树状数组
    for (int i = 1; i <= n; i++) {
        for (int j = i; j <= n; j += j & -j) {
            ft->tree[j] += ft->arr[i];
        }
    }
    
    return ft;
}

// 单点更新：将第i个元素增加delta
void update(FenwickTree *ft, int i, int delta) {
    ft->arr[i] += delta;
    for (; i <= ft->size; i += i & -i) {
        ft->tree[i] += delta;
    }
}

// 查询前缀和：求前i个元素的和
int queryPrefixSum(FenwickTree *ft, int i) {
    int sum = 0;
    for (; i > 0; i -= i & -i) {
        sum += ft->tree[i];
    }
    return sum;
}

// 释放树状数组
void freeFenwickTree(FenwickTree *ft) {
    free(ft->arr);
    free(ft->tree);
    free(ft);
}

int main() {
    int n;
    scanf("%d", &n);
    
    int *data = (int*)malloc(n * sizeof(int));
    for (int i = 0; i < n; i++) {
        scanf("%d", &data[i]);
    }
    
    // 初始化树状数组
    FenwickTree *ft = initFenwickTree(data, n);
    
    // 输出树状数组元素（1-based到n）
    for (int i = 1; i <= n; i++) {
        printf("%d ", ft->tree[i]);
    }
    printf("\n");
    
    // 输出前缀和数组
    for (int i = 1; i <= n; i++) {
        printf("%d ", queryPrefixSum(ft, i));
    }
    printf("\n");
    
    return 0;
}