poj 3321

Intro

题目地址 http://poj.org/problem?id=3321， 我的遍历算法提交上去居然 TLE 了。poj 网站也没有告诉我用了多少秒，这样我就没办法知道我和别人的差距。也找不到测试数据集来自测，只好写个 python 程序来生成测试数据集。

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import Queue
import random  

"""
print 3
print "1 2"
print "1 3"
print "3" 
print "Q 1"
print "C 2"
print "Q 1"
"""


MAX_FORK = 100000
MAX_QUERY = 100000



def build_balanced_fork():
    # 这样建的树是 平衡树
    q = Queue.Queue()  
    q.put(1)
    fork_num = 0
    min_b = 2
    
    while fork_num < MAX_FORK:
        s = q.get()
        print "%d %d
%d %d" % (s, min_b, s, min_b+1)
        fork_num += 2
        q.put(min_b)
        q.put(min_b+1)
        min_b = min_b + 2 

        
def build_unbalanced_fork():
    q = Queue.Queue()  
    q.put(1)
    fork_num = 0
    min_b = 2
    
    while fork_num < MAX_FORK:
        s = q.get()
        if q.qsize() and random.choice([True, False]):
            continue
        print "%d %d
%d %d" % (s, min_b, s, min_b+1)
        fork_num += 2
        q.put(min_b)
        q.put(min_b+1)
        min_b = min_b + 2 

        
        
def build_query():

    for i in range(0, MAX_QUERY):
        operation = random.choice(['C', 'Q'])
        number = random.randint(1, MAX_QUERY)
        print "%s %d" % (operation, number)


if __name__ == '__main__':
    print MAX_FORK+1
    build_unbalanced_fork()
    print MAX_QUERY
    build_query()

这里的 markdown 编辑器没有预览功能，好难用啊！

一开始， 我只实现了 build_balanced_fork， build_query 两个函数。 build_balanced_fork 函数建立的是个平衡树。 在平衡树树上， 我的遍历修改查询程序和采用树状数组的程序（用的是 hzwer 的代码 http://hzwer.com/8114.html ）所花费时间差不多。

用 build_unbalanced_fork 建非平衡树测试集。在非平衡树上，我的程序所用的时间是 hzwer 程序的 几十倍。经过修改后， hzwer 的程序还是比我的快 3 倍。考虑到他的程序用了位运算，而我的程序没有使用。这也是他程序比我快的一个原因。
这样看来， 采用树状数组会减少程序运行时间。

# include <vector>
# include <iostream>
# include <memory.h>


using namespace std;

long long count = 0;

const int maxn = 100010;
int bit[maxn];

int tstart[maxn] = {0};
int tend[maxn] = {0};
int n = 0;

vector<vector<int>> tree(maxn+1);

int lowbit(int x) {
  return x&-x;
}

int sum(int i) {
  int res = 0;
  while (i > 0) {
    res += bit[i];
    i -= lowbit(i);
  }
  return res;
}

void add(int i, int x) {
  while (i <= n) {
    bit[i] += x;
    i += lowbit(i);
  }
}



int build_node(int i) {
  tstart[i] = count;
  count++;
  add(count, 1);
  for (int j=0; j< tree[i].size(); j++) {
    build_node(tree[i][j]);
  }
  tend[i] = count;
  return 0; 
}
               

int main(int argc, char *argv[]) {
  int m = 0;
  int m_index = 0;
  int delta = 0;
  int parent, son;
  long long op = 0;

  scanf("%d", &n);
  vector<int> ans;
  vector<int> has_apple(n+1, 1);

  memset(bit, 0, sizeof(int)*maxn);  // 一开始没有 memset 初始化，导致程结果总是不对！
  
  for (int i=1; i<n; i++) {
    scanf("%d%d", &parent, &son);
    tree[parent].push_back(son);
  }
  
  build_node(1);
  
  scanf("%d", &m);

  while (m--) {
    char cmd[2];
    int now_p = 0;
    
    scanf("%s%d", cmd, &m_index);
    if (cmd[0]== 'C') {
      delta = has_apple[m_index]? -1: 1;
      has_apple[m_index] ^=1;
      add(tstart[m_index]+1, delta);
      }
      
    else {
      ;
      #ifndef DEBUG
      ans.push_back(sum(tend[m_index]) - sum(tstart[m_index]));
     #endif
    }
  }
  #ifdef DEBUG
  printf("%lld
", op);
  #else
  for (int i=0; i< ans.size(); i++)
    printf("%d
", ans[i]);
  #endif
}

经过我不懈的努力，总算用树状数组把时间降下来了。话说位运算就是快。现在来分析一下为什么之前的程序会 TLE 到底性能瓶颈在哪里？

# include <vector>
# include <iostream>


using namespace std;

vector <long long> sum_node (100010, 0);


long long build_node(int i, vector<vector<int>> &tree) {
  long long ans = 1;
  
  for (int j = 0; j < tree[i].size(); j++)
    ans += build_node(tree[i][j], tree);
  
  sum_node[i] = ans;
  return ans;

}
               

int main(int argc, char *argv[]) {
  int n = 0, m = 0;
  int parent, son;
  int delta = 0;
  int deltas = [1, -1];

  long long op = 0;
  scanf("%d", &n);
  vector<long long> ans;
  vector<int> has_apple(n+1, 1);
  vector<int> father_node(n+1, -1);    
  vector<vector <int>> tree(n+1);
  
  for (int i=1; i<n; i++) {
    scanf("%d%d", &parent, &son);
    father_node[son] = parent;    
    tree[parent].push_back(son);
  }
  
  build_node(1, tree);
  
  scanf("%d", &m);

  while (m--) {
    char cmd[2];
    int now_p = 0;
    
    scanf("%s%d", cmd, &n);
    if (cmd[0]== 'C') {
      delta = has_apple[n] > 0? -1: 1;  // 三元操作符要比 delta = deltas[ has_apple[n]] 快，就这个一个操作减少数百（数十？）毫秒（总共也就跑了 1.2 s）
      has_apple[n] = (has_apple[n] + 1)%2;
      now_p = n;
      sum_node[n] += delta;
      
      while (true) {    // 这个更新操作很耗费时间，同样的更新操作。它要比树状数组多两个量级（真是看不懂）
        if (father_node[now_p] == -1)
          break;
        sum_node[father_node[now_p]] += delta;
        now_p = father_node[now_p];    
        #ifdef DEBUG
        op++;
        #endif
      }
      
    }
    else {
      ;
      #ifndef DEBUG
      ans.push_back(sum_node[n]);
     #endif
    }
  }
  #ifdef DEBUG
  printf("%lld
", op);
  #else
  for (int i=0; i< ans.size(); i++)
    printf("%lld
", ans[i]);
  #endif
}

我提交的程序比 hzw 程序在时间上差不多，但是内存用了 8000 k。 不知道线段树能不能过这题？我的代码还是不够优美，网友 bigbigship 的代码很优美！

    #include <cstring>  
    #include <algorithm>  
    #include <cstdio>  
    #include <vector>  
    using namespace std;  
      
    const int maxn = 1e5+10;  //这优雅
      
    int has[maxn];  
    int n;  
      
    struct Tree{  
        struct nod{  
            int to,next;  
        }edge[maxn];  
        int in[maxn],out[maxn],tt;  
        int head[maxn],ip;  
        void init(){  
            tt=0,ip=0;  
            memset(head,-1,sizeof(head));  
        }  
        void add(int u,int v){  
            edge[ip].to= v;  
            edge[ip].next = head[u];  
            head[u]= ip++;  
        }  
        void dfs(int u){  
            in[u]=++tt;  
            for(int i=head[u];i!=-1;i=edge[i].next){  
                int v = edge[i].to;  
                dfs(v);  
            }  
            out[u]=tt;  
        }  
    }G;  
      
    struct BIT {  // 将 bit 变成结构体，优雅
        int sum[maxn];  
        void init(){  
            memset(sum,0,sizeof(sum));  
        }  
        int lowbit(int x) {  
            return x&(-x);  
        }  
        void update(int pos,int x){  
            while(pos<=n){  
                sum[pos]+=x;  
                pos+=lowbit(pos);  
            }  
        }  
        int query(int pos){  
            int ans = 0;  
            while(pos>0){  
                ans+=sum[pos];  
                pos-=lowbit(pos);  
            }  
            return ans;  
        }  
    }T;  
      
    int main() {  
        while(~scanf("%d",&n)){  
            G.init();  
            T.init();  
            for(int i=1;i<n;i++){  
                int u,v;  
                scanf("%d%d",&u,&v);  
                G.add(u,v);  
            }  
            G.dfs(1);  
            for(int i=1;i<=n;i++){  
                T.update(G.in[i],1);  
                has[i]=1;  
            }  
            int x,m;  
            scanf("%d",&m);  
            char s[2];  
            while(m--){  
                scanf("%s%d",s,&x);  
                if(s[0]=='C'){  
                    if(has[x])  
                        T.update(G.in[x],-1);  
                    else  
                        T.update(G.in[x],1);  
                    has[x]^=1;  
                }  
                else  
                    printf("%d
",T.query(G.out[x])-T.query(G.in[x]-1));  
            }  
        }  
        return 0;  
    }  

// 代码很优美

这题用线段树应该不好做，因为线段树是 2 叉平衡树！题目上没有保证建立的树是二叉平衡树！